Skip to content

Commit 743c559

Browse files
fix: handle non-standard array indexes in codegen
1 parent db7780b commit 743c559

File tree

1 file changed

+11
-7
lines changed

1 file changed

+11
-7
lines changed

src/systems/codegen_utils.jl

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ function array_variable_assignments(args...)
3434
# get and/or construct the buffer storing indexes
3535
idxbuffer = get!(
3636
() -> map(Returns((0, 0)), eachindex(arrvar)), var_to_arridxs, arrvar)
37-
idxbuffer[arguments(var)[2:end]...] = (i, j)
37+
Origin(first.(axes(arrvar))...)(idxbuffer)[arguments(var)[2:end]...] = (i, j)
3838
end
3939
end
4040

@@ -59,18 +59,22 @@ function array_variable_assignments(args...)
5959
idxs = SArray{Tuple{size(idxs)...}}(idxs)
6060
end
6161
# view and reshape
62-
push!(assignments,
63-
arrvar
64-
term(reshape, term(view, generated_argument_name(buffer_idx), idxs),
65-
size(arrvar)))
62+
63+
expr = term(reshape, term(view, generated_argument_name(buffer_idx), idxs),
64+
size(arrvar))
6665
else
6766
elems = map(idxs) do idx
6867
i, j = idx
6968
term(getindex, generated_argument_name(i), j)
7069
end
71-
# use `MakeArray` and generate a stack-allocated array
72-
push!(assignments, arrvar MakeArray(elems, SArray))
70+
# use `MakeArray` syntax and generate a stack-allocated array
71+
expr = term(SymbolicUtils.Code.create_array, SArray, nothing,
72+
Val(ndims(arrvar)), Val(length(arrvar)), elems...)
73+
end
74+
if any(x -> !isone(first(x)), axes(arrvar))
75+
expr = term(Origin(first.(axes(arrvar))...), expr)
7376
end
77+
push!(assignments, arrvar expr)
7478
end
7579

7680
return assignments

0 commit comments

Comments
 (0)