ND array transpilation: zeros, ones, indexing, size (column-major flat)

Representation: flat JS Array + _size metadata (column-major, matching Julia).
A[i,j] → A[(j-1)*nrows + (i-1)] — same stride math as WASM compilers.

Construction:
- zeros(m,n), ones(m,n), fill(v,m,n) → jl_ndarray runtime helper
- Creates flat Array with _size property for shape metadata

Indexing:
- A[i,j] read → column-major flat index via _size[0]
- A[i,j] = val → same stride math for write
- A[i,j,k] read → 3D column-major stride
- Handlers in ALL three code paths: compile_invoke, compile_call SSA, compile_call GlobalRef

Size:
- size(A) → A._size.slice() (or [A.length] for 1D)
- size(A,d) → A._size[d-1]

Runtime helper:
- jl_ndarray(fill_val, dims) — creates flat array with _size

Tests: 8 e2e Node.js tests — zeros, ones, length, size(1), size(2),
set+get, column-major verification, unrolled matmul.

Known limitation: nested for loops (for i; for j; for k) have a
pre-existing JST bug where inner break exits outer while loop.
Single for loops work correctly. Filed as TODO.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Dale-Black

src/compiler/codegen.jl

@@ -878,16 +878,21 @@ function compile_call(ctx::JSCompilationContext, expr::Expr)
                 # Array indexing: arr[i] → arr[i-1]
                 # Range indexing: arr[a:b] → arr.slice(a-1, b)
                 if length(call_args) == 2
-                    idx_arg = args[3]  # The index argument (before compilation)
+                    idx_arg = args[3]
                     idx_type = nothing
                     if idx_arg isa Core.SSAValue
                         idx_type = try ctx.code_info.ssavaluetypes[idx_arg.id] catch; nothing end
                     end
-                    # Check if index is a range (UnitRange) → slice
                     if idx_type !== nothing && idx_type isa DataType && idx_type <: AbstractRange
                         return "$(call_args[1]).slice(($(call_args[2])).start-1,($(call_args[2])).stop)"
                     end
                     return "$(call_args[1])[($(call_args[2])) - 1]"
+                elseif length(call_args) == 3
+                    # A[i,j] → column-major flat index: (j-1)*nrows + (i-1)
+                    return "$(call_args[1])[(($(call_args[3]))-1)*$(call_args[1])._size[0]+(($(call_args[2]))-1)]"
+                elseif length(call_args) == 4
+                    # A[i,j,k] → column-major: (k-1)*m*n + (j-1)*m + (i-1)
+                    return "$(call_args[1])[(($(call_args[4]))-1)*$(call_args[1])._size[0]*$(call_args[1])._size[1]+(($(call_args[3]))-1)*$(call_args[1])._size[0]+(($(call_args[2]))-1)]"
                 end
                 return "[]"
             end
@@ -1031,6 +1036,18 @@ function compile_call(ctx::JSCompilationContext, expr::Expr)
             if fn_name == "sort" && length(call_args) >= 1
                 return "$(call_args[1]).slice().sort()"
             end
+            if fn_name == "setindex!" && length(call_args) >= 3
+                if length(call_args) == 4
+                    # A[i,j] = val → column-major
+                    return "($(call_args[1])[(($(call_args[4]))-1)*$(call_args[1])._size[0]+(($(call_args[3]))-1)] = $(call_args[2]))"
+                else
+                    return "($(call_args[1])[($(call_args[3]))-1] = $(call_args[2]))"
+                end
+            end
+            if fn_name == "getindex" && length(call_args) == 3
+                # A[i,j] → column-major
+                return "$(call_args[1])[(($(call_args[3]))-1)*$(call_args[1])._size[0]+(($(call_args[2]))-1)]"
+            end
             if fn_name == "copy" && length(call_args) >= 1
                 return "$(call_args[1]).slice()"
             end
@@ -1058,15 +1075,39 @@ function compile_call(ctx::JSCompilationContext, expr::Expr)
                 return "$(call_args[2]).reduce($(call_args[1]))"
             end
 
-            # ─── Construction (by name) ───
-            if fn_name == "zeros" && length(call_args) >= 1
-                return "new Array($(call_args[1])).fill(0)"
+            # ─── Construction (by name, 1D and ND) ───
+            if fn_name == "zeros"
+                if length(call_args) == 1
+                    return "new Array($(call_args[1])).fill(0)"
+                elseif length(call_args) >= 2
+                    require_runtime!(ctx, :jl_ndarray)
+                    return "jl_ndarray(0,[$(join(call_args, ","))])"
+                end
             end
-            if fn_name == "ones" && length(call_args) >= 1
-                return "new Array($(call_args[1])).fill(1)"
+            if fn_name == "ones"
+                if length(call_args) == 1
+                    return "new Array($(call_args[1])).fill(1)"
+                elseif length(call_args) >= 2
+                    require_runtime!(ctx, :jl_ndarray)
+                    return "jl_ndarray(1,[$(join(call_args, ","))])"
+                end
             end
-            if fn_name == "fill" && length(call_args) >= 2
-                return "new Array($(call_args[2])).fill($(call_args[1]))"
+            if fn_name == "fill"
+                if length(call_args) == 2
+                    return "new Array($(call_args[2])).fill($(call_args[1]))"
+                elseif length(call_args) >= 3
+                    require_runtime!(ctx, :jl_ndarray)
+                    return "jl_ndarray($(call_args[1]),[$(join(call_args[2:end], ","))])"
+                end
+            end
+
+            # ─── Size (by name) ───
+            if fn_name == "size"
+                if length(call_args) == 1
+                    return "($(call_args[1])._size||[$(call_args[1]).length]).slice()"
+                elseif length(call_args) == 2
+                    return "($(call_args[1])._size?$(call_args[1])._size[($(call_args[2]))-1]:$(call_args[1]).length)"
+                end
             end
 
             # ─── Parsing (by name) ───
@@ -1265,10 +1306,24 @@ function compile_call(ctx::JSCompilationContext, expr::Expr)
                     return "[]"
                 end
             end
-            # Array indexing
+            # Array indexing (1D and ND)
             call_args_gr = [compile_value(ctx, a) for a in args[2:end]]
             if length(call_args_gr) == 2
                 return "$(call_args_gr[1])[($(call_args_gr[2])) - 1]"
+            elseif length(call_args_gr) == 3
+                # A[i,j] → column-major
+                return "$(call_args_gr[1])[(($(call_args_gr[3]))-1)*$(call_args_gr[1])._size[0]+(($(call_args_gr[2]))-1)]"
+            end
+        end
+
+        # Base.setindex! — array assignment (1D and ND)
+        if bname === :setindex! && callee.mod === Base
+            call_args_gr = [compile_value(ctx, a) for a in args[2:end]]
+            if length(call_args_gr) == 4
+                # A[i,j] = val → column-major
+                return "($(call_args_gr[1])[(($(call_args_gr[4]))-1)*$(call_args_gr[1])._size[0]+(($(call_args_gr[3]))-1)] = $(call_args_gr[2]))"
+            elseif length(call_args_gr) == 3
+                return "($(call_args_gr[1])[($(call_args_gr[3]))-1] = $(call_args_gr[2]))"
             end
         end
 
@@ -1677,8 +1732,15 @@ function compile_invoke(ctx::JSCompilationContext, expr::Expr)
             elseif func_name == "setindex!"
                 arr_val = compile_value(ctx, expr.args[3])
                 val_val = compile_value(ctx, expr.args[4])
-                idx_val = compile_value(ctx, expr.args[5])
-                return "($(arr_val)[($(idx_val))-1] = $(val_val))"
+                if length(expr.args) == 6
+                    # A[i,j] = val → column-major
+                    i_val = compile_value(ctx, expr.args[5])
+                    j_val = compile_value(ctx, expr.args[6])
+                    return "($(arr_val)[(($(j_val))-1)*$(arr_val)._size[0]+(($(i_val))-1)] = $(val_val))"
+                else
+                    idx_val = compile_value(ctx, expr.args[5])
+                    return "($(arr_val)[($(idx_val))-1] = $(val_val))"
+                end
             elseif func_name == "deleteat!"
                 arr_val = compile_value(ctx, expr.args[3])
                 idx_val = compile_value(ctx, expr.args[4])
@@ -1825,6 +1887,15 @@ function compile_invoke(ctx::JSCompilationContext, expr::Expr)
         return "$(call_args[1]).length === 0"
     end
 
+    # size(A) → shape tuple, size(A, d) → dimension size
+    if func_name == "size"
+        if length(call_args) == 1
+            return "($(call_args[1])._size||[$(call_args[1]).length]).slice()"
+        elseif length(call_args) == 2
+            return "($(call_args[1])._size?$(call_args[1])._size[($(call_args[2]))-1]:$(call_args[1]).length)"
+        end
+    end
+
     # convert(T, x) → just x (type conversions are compile-time in JS)
     if func_name == "convert" && length(call_args) >= 2
         return call_args[2]

src/compiler/runtime.jl

@@ -218,6 +218,15 @@ function jl_objectid(x) {
   if (id === undefined) { id = ++_jl_oid_ctr; _jl_oid_map.set(x, id); }
   return id;
 }""",
+
+    :jl_ndarray => """
+function jl_ndarray(fill_val, dims) {
+  var n = 1;
+  for (var i = 0; i < dims.length; i++) n *= dims[i];
+  var a = new Array(n).fill(fill_val);
+  a._size = dims.slice();
+  return a;
+}""",
 )
 
 # Dependency-ordered list of symbols for deterministic output
@@ -230,6 +239,7 @@ const RUNTIME_ORDER = [
     :jl_checked_add, :jl_checked_sub, :jl_checked_mul,
     :jl_println, :jl_print,
     :jl_objectid,
+    :jl_ndarray,
 ]
 
 """

test/runtests.jl

@@ -3515,4 +3515,43 @@ process.stdout.write(JSON.stringify(r));
         fn = () -> startswith("hello", "he")
         @test compile_unopt_and_run(fn, "") == "true"
     end
+
+    # ─── ND Array tests (e2e via Node.js) ───
+    @testset "ND: zeros(3,4)" begin
+        @test compile_unopt_and_run(() -> zeros(3, 4), "") == "[0,0,0,0,0,0,0,0,0,0,0,0]"
+    end
+    @testset "ND: ones(2,3)" begin
+        @test compile_unopt_and_run(() -> ones(2, 3), "") == "[1,1,1,1,1,1]"
+    end
+    @testset "ND: length(zeros(3,4))" begin
+        @test compile_unopt_and_run(() -> length(zeros(3, 4)), "") == "12"
+    end
+    @testset "ND: size(A,1)" begin
+        @test compile_unopt_and_run(() -> size(zeros(3, 4), 1), "") == "3"
+    end
+    @testset "ND: size(A,2)" begin
+        @test compile_unopt_and_run(() -> size(zeros(3, 4), 2), "") == "4"
+    end
+    @testset "ND: A[i,j] set+get" begin
+        @test compile_unopt_and_run(() -> begin A=zeros(2,3); A[1,2]=42.0; A[1,2] end, "") == "42"
+    end
+    @testset "ND: column-major indexing" begin
+        @test compile_unopt_and_run(() -> begin A=zeros(2,3); A[2,1]=10.0; A[1,2]=20.0; A[2,2]=30.0; A end, "") == "[0,10,20,30,0,0]"
+    end
+    @testset "ND: unrolled matmul" begin
+        fn = () -> begin
+            A=zeros(2,2); A[1,1]=1.0; A[1,2]=2.0; A[2,1]=3.0; A[2,2]=4.0
+            B=zeros(2,2); B[1,1]=5.0; B[1,2]=6.0; B[2,1]=7.0; B[2,2]=8.0
+            C=zeros(2,2)
+            C[1,1] = A[1,1]*B[1,1] + A[1,2]*B[2,1]
+            C[2,1] = A[2,1]*B[1,1] + A[2,2]*B[2,1]
+            C[1,2] = A[1,1]*B[1,2] + A[1,2]*B[2,2]
+            C[2,2] = A[2,1]*B[1,2] + A[2,2]*B[2,2]
+            return C
+        end
+        @test compile_unopt_and_run(fn, "") == "[19,43,22,50]"
+    end
+    # Note: nested for loops (for i; for j; for k) have a known JST bug
+    # where inner loop `break` exits the outer while loop. Single for loops work.
+    # TODO: fix nested for loop compilation in JST
 end