docs: add docstring for code generation utils

Author: AayushSabharwal

src/systems/codegen_utils.jl

@@ -1,66 +1,134 @@
+"""
+    $(TYPEDSIGNATURES)
+
+Return the name for the `i`th argument in a function generated by `build_function_wrapper`.
+"""
 function generated_argument_name(i::Int)
     return Symbol(:__mtk_arg_, i)
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Given the arguments to `build_function_wrapper`, return a list of assignments which
+reconstruct array variables if they are present scalarized in `args`.
+"""
 function array_variable_assignments(args...)
+    # map array symbolic to an identically sized array where each element is (buffer_idx, idx_in_buffer)
     var_to_arridxs = Dict{BasicSymbolic, Array{Tuple{Int, Int}}}()
     for (i, arg) in enumerate(args)
+        # filter out non-arrays
+        # any element of args which is not an array is assumed to not contain a
+        # scalarized array symbolic. This works because the only non-array element
+        # is the independent variable
         symbolic_type(arg) == NotSymbolic() || continue
         arg isa AbstractArray || continue
 
+        # go through symbolics
         for (j, var) in enumerate(arg)
             var = unwrap(var)
+            # filter out non-array-symbolics
             iscall(var) || continue
             operation(var) == getindex || continue
             arrvar = arguments(var)[1]
+            # get and/or construct the buffer storing indexes
             idxbuffer = get!(() -> map(Returns((0, 0)), eachindex(arrvar)), var_to_arridxs, arrvar)
             idxbuffer[arguments(var)[2:end]...] = (i, j)
         end
     end
 
     assignments = Assignment[]
     for (arrvar, idxs) in var_to_arridxs
+        # all elements of the array need to be present in `args` to form the
+        # reconstructing assignment
         any(iszero ∘ first, idxs) && continue
 
+        # if they are all in the same buffer, we can take a shortcut and `view` into it
         if allequal(Iterators.map(first, idxs))
             buffer_idx = first(first(idxs))
             idxs = map(last, idxs)
+            # if all the elements are contiguous and ordered, turn the array of indexes into a range
+            # to help reduce allocations
             if first(idxs) < last(idxs) && vec(idxs) == first(idxs):last(idxs)
                 idxs = first(idxs):last(idxs)
             elseif vec(idxs) == last(idxs):-1:first(idxs)
                 idxs = last(idxs):-1:first(idxs)
             else
+                # Otherwise, turn the indexes into an `SArray` so they're stack-allocated
                 idxs = SArray{Tuple{size(idxs)...}}(idxs)
             end
+            # view and reshape
             push!(assignments, arrvar ← term(reshape, term(view, generated_argument_name(buffer_idx), idxs), size(arrvar)))
         else
             elems = map(idxs) do idx
                 i, j = idx
                 term(getindex, generated_argument_name(i), j)
             end
+            # use `MakeArray` and generate a stack-allocated array
             push!(assignments, arrvar ← MakeArray(elems, SArray))
         end
     end
 
     return assignments
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+A wrapper around `build_function` which performs the necessary transformations for
+code generation of all types of systems. `expr` is the expression returned from the
+generated functions, and `args` are the arguments.
+
+# Keyword Arguments
+
+- `p_start`, `p_end`: Denotes the indexes in `args` where the buffers of the splatted
+  `MTKParameters` object are present. These are collapsed into a single argument and
+  destructured inside the function. `p_start` must also be provided for non-split systems
+  since it is used by `wrap_delays`.
+- `wrap_delays`: Whether to transform delayed unknowns of `sys` present in `expr` into
+  calls to a history function. The history function is added to the list of arguments
+  right before parameters, at the index `p_start`.
+- `wrap_code`: Forwarded to `build_function`.
+- `add_observed`: Whether to add assignment statements for observed equations in the
+  generated code.
+- `filter_observed`: A predicate function to filter out observed equations which should
+  not be added to the generated code.
+- `create_bindings`: Whether to explicitly destructure arrays of symbolics present in
+  `args` in the generated code. If `false`, all usages of the individual symbolics will
+  instead call `getindex` on the relevant argument. This is useful if the generated
+  function writes to one of its arguments and expects subsequent code to use the new
+  values. Note that the collapsed `MTKParameters` argument will always be explicitly
+  destructured regardless of this keyword argument.
+- `output_type`: The type of the output buffer. If `mkarray` (see below) is `nothing`,
+  this will be passed to the `similarto` argument of `build_function`. If `output_type`
+  is `Tuple`, `expr` will be wrapped in `SymbolicUtils.Code.MakeTuple` (regardless of
+  whether it is scalar or an array).
+- `mkarray`: A function which accepts `expr` and `output_type` and returns a code
+  generation object similar to `MakeArray` or `MakeTuple` to be used to generate
+  code for `expr`.
+- `wrap_mtkparameters`: Whether to collapse parameter buffers for a split system into a
+  argument.
+
+All other keyword arguments are forwarded to `build_function`.
+"""
 function build_function_wrapper(sys::AbstractSystem, expr, args...; p_start = 2, p_end = is_time_dependent(sys) ? length(args) - 1 : length(args), wrap_delays = is_dde(sys), wrap_code = identity, add_observed = true, filter_observed = Returns(true), create_bindings = true, output_type = nothing, mkarray = nothing, wrap_mtkparameters = true, kwargs...)
     isscalar = !(expr isa AbstractArray || symbolic_type(expr) == ArraySymbolic())
-
+    # filter observed equations
     obs = filter(filter_observed, observed(sys))
+    # turn delayed unknowns into calls to the history function
     if wrap_delays
         history_arg = is_split(sys) ? MTKPARAMETERS_ARG : generated_argument_name(p_start)
         obs = map(obs) do eq
             delay_to_function(sys, eq; history_arg)
         end
         expr = delay_to_function(sys, expr; history_arg)
+        # add extra argument
         args = (args[1:p_start-1]..., DDE_HISTORY_FUN, args[p_start:end]...)
         p_start += 1
         p_end += 1
     end
     pdeps = parameter_dependencies(sys)
-
+    # get the constants to add to the code
     cmap, _ = get_cmap(sys)
     extra_constants = collect_constants(expr)
     filter!(extra_constants) do c
@@ -69,13 +137,15 @@ function build_function_wrapper(sys::AbstractSystem, expr, args...; p_start = 2,
     for c in extra_constants
         push!(cmap, c ~ getdefault(c))
     end
+    # only get the necessary observed equations, avoiding extra computation
     if add_observed
         obsidxs = observed_equations_used_by(sys, expr)
     else
         obsidxs = Int[]
     end
+    # similarly for parameter dependency equations
     pdepidxs = observed_equations_used_by(sys, expr; obs = pdeps)
-
+    # assignments for reconstructing scalarized array symbolics
     assignments = array_variable_assignments(args...)
 
     for eq in Iterators.flatten((cmap, pdeps[pdepidxs], obs[obsidxs]))
@@ -84,35 +154,43 @@ function build_function_wrapper(sys::AbstractSystem, expr, args...; p_start = 2,
 
     args = ntuple(Val(length(args))) do i
         arg = args[i]
+        # for time-dependent systems, all arguments are passed through `time_varying_as_func`
+        # TODO: This is legacy behavior and a candidate for removal in v10 since we have callable
+        # parameters now.
         if is_time_dependent(sys)
             arg = if symbolic_type(arg) == NotSymbolic()
                 arg isa AbstractArray ? map(x -> time_varying_as_func(unwrap(x), sys), arg) : arg
             else
                 time_varying_as_func(unwrap(arg), sys)
             end
         end
+        # Make sure to use the proper names for arguments
         if symbolic_type(arg) == NotSymbolic() && arg isa AbstractArray
             DestructuredArgs(arg, generated_argument_name(i); create_bindings)
         else
             arg
         end
     end
 
+    # wrap into a single MTKParameters argument
     if is_split(sys) && wrap_mtkparameters
         if p_start > p_end
+            # In case there are no parameter buffers, still insert an argument
             args = (args[1:p_start-1]..., MTKPARAMETERS_ARG, args[p_end+1:end]...)
         else
             # cannot apply `create_bindings` here since it doesn't nest
             args = (args[1:p_start-1]..., DestructuredArgs(collect(args[p_start:p_end]), MTKPARAMETERS_ARG), args[p_end+1:end]...)
         end
     end
 
+    # add preface assignments
     if has_preface(sys) && (pref = preface(sys)) !== nothing
         append!(assignments, pref)
     end
 
     wrap_code = wrap_code .∘ wrap_assignments(isscalar, assignments)
 
+    # handling of `output_type` and `mkarray`
     similarto = nothing
     if output_type === Tuple
         expr = MakeTuple(Tuple(expr))
@@ -124,6 +202,7 @@ function build_function_wrapper(sys::AbstractSystem, expr, args...; p_start = 2,
         wrap_code = wrap_code[2]
     end
 
+    # scalar `build_function` only accepts a single function for `wrap_code`.
     if wrap_code isa Tuple && symbolic_type(expr) == ScalarSymbolic()
         wrap_code = wrap_code[1]
     end