Need help with sparse matrix for loop code

I'm working on updating a branch of the compiler that supports sparse matrices and having trouble writing the function that makes a for loop over the sparse matrix to fill it with data. There I had a function that generated the for loop for reading in data like link

let mkfortnite nonzero_rows nonzero_cols bodyfn var smeta =
  let idx s =
    Single {expr= Var s; emeta= {mtype= UInt; mloc = smeta; madlevel= DataOnly}}
  in
  let loopvar, reset = gensym_enter () in
  let lower = loop_bottom in
  let upper = internal_funapp FnLength [nonzero_rows] internal_meta in
  let row_idx = add_int_index nonzero_rows (idx loopvar) in
  let col_idx = add_int_index nonzero_cols (idx loopvar) in
  let bodyfn var =
    add_int_index (add_int_index var (Single row_idx)) (Single col_idx)
    |> bodyfn in
  reset ();
  {stmt= For {loopvar; lower; upper; body=bodyfn var}; smeta}

Where the old version looked like

(** [mkfor] returns a MIR For statement that iterates over the given expression
    [iteratee]. *)
let mkfor upper bodyfn iteratee smeta =
  let idx s =
    Single {expr= Var s; emeta= {mtype= UInt; mloc= smeta; madlevel= DataOnly}}
  in
  let loopvar, reset = gensym_enter () in
  let lower = loop_bottom in
  let stmt = Block [bodyfn (add_int_index iteratee (idx loopvar))] in
  reset () ;
  {stmt= For {loopvar; lower; upper; body= {stmt; smeta}}; smeta}

The new version of mkfor looks like

  (** [mkfor] returns a MIR For statement that iterates over the given expression
    [iteratee]. *)
  let mkfor upper bodyfn iteratee meta =
    let idx s =
      let meta =
        Expr.Typed.Meta.create ~type_:UInt ~loc:meta ~adlevel:DataOnly ()
      in
      let expr = Expr.Fixed.{meta; pattern= Var s} in
      Index.Single expr
    in
    let loopvar, reset = Gensym.enter () in
    let lower = Expr.Helpers.loop_bottom in
    let stmt =
      Fixed.Pattern.Block
        [bodyfn (Expr.Helpers.add_int_index iteratee (idx loopvar))]
    in
    reset () ;
    let body = Fixed.{meta; pattern= stmt} in
    let pattern = Fixed.Pattern.For {loopvar; lower; upper; body} in
    Fixed.{meta; pattern}

I think I'm pretty close, I've written the below but nonzero_rows in row_idx and upper in pattern are giving me the errors

row_idx: (nonzero_rows)

type _ = unit Expr.Fixed.t
This expression has type unit Expr.Fixed.t
but an expression was expected of type
  Expr.Typed.t = Expr.Typed.Meta.t Expr.Fixed.t
Type unit is not compatible with type Expr.Typed.Meta.t

pattern: (upper)

type _ = unit Expr.Fixed.t
This expression has type unit Expr.Fixed.t
but an expression was expected of type
  Expr.Typed.t = Expr.Typed.Meta.t Expr.Fixed.t
Type unit is not compatible with type Expr.Typed.Meta.t

  let mkfortnite nonzero_rows nonzero_cols bodyfn var meta = 
    let idx s = 
       let meta = 
         Expr.Typed.Meta.create ~type_:Uint ~loc:meta ~adlevel:DataOnly () 
        in
        let expr = Expr.Fixed.{meta; pattern= Var s} in
        Index.Single expr
    in
    let loopvar, reset = Gensym.enter () in
    let lower = Expr.Helpers.loop_bottom in
    let upper = Expr.Helpers.internal_funapp FnLength [nonzero_rows] () in (* error from here*)
    let row_idx = Expr.Helpers.add_int_index nonzero_rows (idx loopvar) in (* error here*)
    let col_idx = Expr.Helpers.add_int_index nonzero_cols (idx loopvar) in
    let row_iter = Expr.Helpers.add_int_index var (Index.Single row_idx) in
    let col_iter = Expr.Helpers.add_int_index row_iter (Index.Single col_idx) in
    let bodyfn = col_iter |> bodyfn in
    reset ();
    let body = Fixed.{meta; pattern = bodyfn var} in
    let pattern = Fixed.Pattern.For {loopvar; lower; upper ; body} in (* error here*)
    Fixed.{meta; pattern}

mkfortnite is called in for_scalar here

  let rec for_scalar st bodyfn var smeta =
    match st with
    | SizedType.SInt | SReal -> bodyfn var
    | SVector d | SRowVector d -> mkfor d bodyfn var smeta
    | SMatrix (d1, d2) ->
        mkfor d1 (fun e -> for_scalar (SRowVector d2) bodyfn e smeta) var smeta
    | SSparseMatrix (nonzero_rows, nonzero_cols, _, _) ->
      mkfortnite nonzero_rows nonzero_cols bodyfn var smeta
    | SArray (t, d) -> mkfor d (fun e -> for_scalar t bodyfn e smeta) var smeta

The goal here is to generate a for loop like

      x = Eigen::SparseMatrix<double>(N, M);
      current_statement__ = 7;
      pos__ = 1;
      current_statement__ = 7;
// nz_rows has a size equal to the number of nonzero values
      for (size_t sym1__ = 1; sym1__ <= stan::math::size(nz_rows); ++sym1__) {
        current_statement__ = 7;
        assign(x,
               cons_list(index_uni(nz_cols[(sym1__ - 1)]),
                         cons_list(index_uni(nz_rows[(sym1__ - 1)]),
                                   nil_index_list())),
               context__.vals_r("x")[(pos__ - 1)], "assigning variable x");
        current_statement__ = 7;
        pos__ = (pos__ + 1);
      }

The main things are

1. I'm not totally sure how to use the new types when making upper where it used to be

  let upper = internal_funapp FnLength [nonzero_rows] internal_meta in

I tried

let upper = Expr.Helpers.internal_funapp FnLength [nonzero_rows] () in

It looks like @enetsee removed internal_meta in 3f037c5be9760be14c4e4b2d41cc4dd560a552e0. @enetsee can you help me find the right setup for this now?

2. upper above is causing ocaml to think nonzero_rows is unit Expr.Fixed.t when it should be the same as nonzero_cols (Expr.Typed.t). So I think the heart of the problem is the upper statement.

If anyone can help with this it would be v appreciated! You can grab the branch via

git remote add stevebronder  https://github.com/SteveBronder/stanc3
git fetch stevebronder
git checkout stevebronder/sparse3
make

Start with the mkfortnite and for_scalar entry points on the referenced sparse-matrix branch, then compare their typed-expression construction with the newer mkfor implementation. Build the sparse3 branch and verify that the type errors are resolved and that the generated code produces the shown sparse-matrix loop.