fast path for iterating through top level symbols

2 files changed, 103 insertions, 13 deletions

diff --git a/src/js_ast.zig b/src/js_ast.zig
index 43e57a2f8..27dab5b1a 100644
--- a/src/js_ast.zig
+++ b/src/js_ast.zig
@@ -875,6 +875,21 @@ pub const Symbol = struct {
         // "Ref" for more detail.
         symbols_for_source: NestedList = NestedList{},
 
+        pub fn assignChunkIndex(this: *Map, decls_: DeclaredSymbol.List, chunk_index: u32) void {
+            const Iterator = struct {
+                map: *Map,
+                chunk_index: u32,
+
+                pub fn next(self: @This(), ref: Ref) void {
+                    var symbol = self.map.get(ref).?;
+                    symbol.chunk_index = self.chunk_index;
+                }
+            };
+            var decls = decls_;
+
+            DeclaredSymbol.forEachTopLevelSymbol(&decls, Iterator{ .map = this, .chunk_index = chunk_index }, Iterator.next);
+        }
+
         pub fn merge(this: *Map, old: Ref, new: Ref) Ref {
             if (old.eql(new)) {
                 return new;
@@ -4731,6 +4746,73 @@ pub const DeclaredSymbol = struct {
     is_top_level: bool = false,
 
     pub const List = bun.MultiArrayList(DeclaredSymbol);
+
+    pub fn forEachTopLevelSymbol(decls_: *List, ctx: anytype, comptime Fn: anytype) void {
+        const FnType = @TypeOf(Fn);
+        const ReturnType = bun.meta.ReturnOfType(FnType);
+        var decls = decls_.slice();
+        var is_top_levels_bool = decls.items(.is_top_level);
+        var refs = decls.items(.ref);
+
+        if (comptime Environment.enableSIMD) {
+            const vector_size = 16;
+            if (refs.len >= vector_size) {
+                const refs_end = refs.ptr + (refs.len - refs.len % vector_size);
+                while (refs.ptr != refs_end) {
+                    const vec = @as(@Vector(vector_size, bool), @bitCast([vector_size]bool, is_top_levels_bool.ptr[0..vector_size]));
+
+                    if (@reduce(.Max, vec) > 0) {
+                        var j = 0;
+                        while (j < vector_size) : (j += 1) {
+                            if (vec[j]) {
+                                if (comptime ReturnType == void) {
+                                    Fn(ctx, refs[j]);
+                                } else if (Fn(ctx, refs[j])) |result| {
+                                    refs[j] = result;
+                                }
+                            }
+                        }
+                    }
+
+                    is_top_levels_bool = is_top_levels_bool[vector_size..];
+                    refs = refs[vector_size..];
+                }
+            }
+        }
+
+        const scalar = 8;
+        while (refs.len >= scalar) : (refs = refs[scalar..]) {
+            const eight: u64 = @as(u64, is_top_levels_bool[0..scalar].*);
+            if (eight == 0) {
+                is_top_levels_bool = is_top_levels_bool[scalar..];
+                continue;
+            }
+
+            comptime var j: usize = 0;
+
+            inline while (j < scalar) : (j += 1) {
+                if (is_top_levels_bool[j]) {
+                    if (comptime ReturnType == void) {
+                        Fn(ctx, refs[j]);
+                    } else if (Fn(ctx, refs[j])) |result| {
+                        refs[j] = result;
+                    }
+                }
+            }
+
+            is_top_levels_bool = is_top_levels_bool[scalar..];
+        }
+
+        while (refs.len > 0) : (refs = refs[1..]) {
+            if (comptime ReturnType == void) {
+                Fn(ctx, refs[0]);
+            } else if (Fn(ctx, refs[0])) |result| {
+                refs[0] = result;
+            }
+
+            is_top_levels_bool = is_top_levels_bool[1..];
+        }
+    }
 };
 
 pub const Dependency = struct {
diff --git a/src/js_parser.zig b/src/js_parser.zig
index 8b76a5c1d..7945f48e5 100644
--- a/src/js_parser.zig
+++ b/src/js_parser.zig
@@ -120,7 +120,9 @@ fn foldStringAddition(lhs: Expr, rhs: Expr) ?Expr {
             if (rhs.data == .e_string and left.isUTF8() and rhs.data.e_string.isUTF8()) {
                 var orig = lhs.data.e_string.*;
                 const rhs_clone = Expr.init(E.String, rhs.data.e_string.*, rhs.loc);
-                orig.push(rhs_clone.data.e_string);
+                orig.push(
+                    rhs_clone.data.e_string,
+                );
 
                 return Expr.init(E.String, orig, lhs.loc);
             }
@@ -18803,28 +18805,34 @@ fn NewParser_(
                 var i: u32 = 0;
                 const count = @truncate(u32, parts.len);
                 while (i < count) : (i += 1) {
-                    const decls = parts[i].declared_symbols;
-                    var is_top_level_slice = decls.items(.is_top_level);
-                    var refs = decls.items(.ref);
-                    // TODO: SIMD this loop
-                    for (is_top_level_slice) |is_top_level, j| {
-                        if (is_top_level) {
+                    var decls = &parts[i].declared_symbols;
+                    var ctx_ = .{
+                        .allocator = p.allocator,
+                        .top_level_symbols_to_parts = &top_level_symbols_to_parts,
+                        .symbols = p.symbols.items,
+                        .part_index = i,
+                    };
+                    const Ctx = @TypeOf(ctx_);
+                    const Iterator = struct {
+                        pub fn next(ctx: Ctx, input: Ref) void {
                             // If this symbol was merged, use the symbol at the end of the
                             // linked list in the map. This is the case for multiple "var"
                             // declarations with the same name, for example.
-                            var ref = refs[j];
-                            while (p.symbols.items[ref.innerIndex()].link.isValid()) {
-                                ref = p.symbols.items[ref.innerIndex()].link;
+                            var ref = input;
+                            while (ctx.symbols[ref.innerIndex()].link.isValid()) {
+                                ref = ctx.symbols[ref.innerIndex()].link;
                             }
 
-                            var entry = top_level_symbols_to_parts.getOrPut(p.allocator, ref) catch unreachable;
+                            var entry = ctx.top_level_symbols_to_parts.getOrPut(ctx.allocator, ref) catch unreachable;
                             if (!entry.found_existing) {
                                 entry.value_ptr.* = .{};
                             }
 
-                            entry.value_ptr.push(p.allocator, @truncate(u32, i)) catch unreachable;
+                            entry.value_ptr.push(ctx.allocator, @truncate(u32, ctx.part_index)) catch unreachable;
                         }
-                    }
+                    };
+
+                    DeclaredSymbol.forEachTopLevelSymbol(&decls, ctx_, Iterator.next);
                 }
 
                 // Pulling in the exports of this module always pulls in the export part