fix(library/init/lean/compiler/pushproj): bug and cleanup

library/init/data/array/basic.lean

@@ -295,7 +295,7 @@ let r : Array α := mkEmpty (e - b) in
 if h : e ≤ a.size then extractAux a b e h r
 else r
 
-@[inline] protected def append (a : Array α) (b : Array α) : Array α :=
+protected def append (a : Array α) (b : Array α) : Array α :=
 b.foldl (λ a v, a.push v) a
 
 instance : HasAppend (Array α) := ⟨Array.append⟩

library/init/lean/compiler/default.lean

@@ -5,4 +5,4 @@ Authors: Leonardo de Moura
 -/
 prelude
 import init.lean.compiler.constfolding
-import init.lean.compiler.ir
+import init.lean.compiler.ir init.lean.compiler.pushproj

library/init/lean/compiler/pushproj.lean

@@ -9,27 +9,29 @@ namespace Lean
 namespace IR
 
 partial def pushProjs : Array FnBody → Array Alt → Array VarIdxSet → Array FnBody → VarIdxSet → Array FnBody × Array Alt
-| bs alts fvs ps psvs :=
+| bs alts afvs ps vs :=
   if bs.isEmpty then (ps.reverse, alts)
   else
-    let b  := bs.back in
-    let bs := bs.pop in
+    let b    := bs.back in
+    let bs   := bs.pop in
+    let done := λ _ : Unit, ((bs.push b) ++ ps.reverse, alts) in
+    let skip := λ _ : Unit, pushProjs bs alts afvs (ps.push b) (b.collectFreeVars vs) in
     match b with
     | FnBody.vdecl x t v _ :=
       (match v with
        | Expr.proj _ _    :=
-         if !psvs.contains x.idx && !fvs.all (λ s, s.contains x.idx) then
+         if !vs.contains x.idx && !afvs.all (λ s, s.contains x.idx) then
            let alts := alts.hmapIdx $ λ i alt, alt.modifyBody $ λ b',
-              if (fvs.get i).contains x.idx then b.setBody b'
+              if (afvs.get i).contains x.idx then b.setBody b'
               else b' in
-           let fvs  := fvs.hmap $ λ s, if s.contains x.idx then b.collectFreeVars s else s in
-           pushProjs bs alts fvs ps psvs
+           let fvs  := afvs.hmap $ λ s, if s.contains x.idx then b.collectFreeVars s else s in
+           pushProjs bs alts fvs ps vs
          else
-           pushProjs bs alts fvs (ps.push b) (b.collectFreeVars psvs)
-       | Expr.uproj _ _   := pushProjs bs alts fvs (ps.push b) (b.collectFreeVars psvs)
-       | Expr.sproj _ _ _ := pushProjs bs alts fvs (ps.push b) (b.collectFreeVars psvs)
-       | other            := ((bs.push b) ++ ps.reverse, alts))
-    | other := ((bs.push b) ++ ps.reverse, alts)
+           skip ()
+       | Expr.uproj _ _   := skip ()
+       | Expr.sproj _ _ _ := skip ()
+       | _                := done ())
+    | _ := done ()
 
 /-- Push projections inside `cases` branches. -/
 partial def FnBody.pushProj : FnBody → FnBody
@@ -37,8 +39,10 @@ partial def FnBody.pushProj : FnBody → FnBody
   let (bs, term) := b.flatten in
   match term with
   | FnBody.case tid x alts :=
-    let fvs        := alts.map $ λ alt, alt.body.freeVars in
-    let (bs, alts) := pushProjs bs alts fvs Array.empty {} in
+    let afvs       := alts.map $ λ alt, alt.body.freeVars in
+    let vs         := ({} : VarIdxSet) in
+    let vs         := vs.insert x.idx in
+    let (bs, alts) := pushProjs bs alts afvs Array.empty vs in
     let alts       := alts.hmap $ λ alt, alt.modifyBody $ λ b, FnBody.pushProj b in
     let term       := FnBody.case tid x alts in
     reshape bs term

src/library/compiler/ir.cpp

@@ -38,6 +38,7 @@ object * mk_alt_core(object * n, object * cidx, object * size, object * usize, o
 object * mk_decl_core(object * f, object * xs, uint8 ty, object * b);
 object * decl_to_string_core(object * d);
 object * decl_max_var_core(object * d);
+object * decl_push_proj_core(object * d);
 /*
 inductive IRType
 | float | uint8 | uint16 | uint32 | uint64 | usize
@@ -118,6 +119,10 @@ nat decl_max_var(decl const & d) {
     inc(d.raw());
     return nat(decl_max_var_core(d.raw()));
 }
+decl decl_push_proj(decl const & d) {
+    inc(d.raw());
+    return decl(decl_push_proj_core(d.raw()));
+}
 }
 
 class to_ir_fn {

src/library/compiler/ir.h

@@ -12,6 +12,7 @@ namespace ir {
 typedef object_ref decl;
 std::string decl_to_string(decl const & d);
 nat decl_max_var(decl const & d);
+decl decl_push_proj(decl const & d);
 }
 ir::decl to_ir_decl(environment const & env, comp_decl const & d);
 }

src/library/compiler/llnf.cpp

@@ -2419,6 +2419,9 @@ static void display_ir(environment const & env, comp_decl const & decl) {
     ir::decl d = to_ir_decl(env, decl);
     tout() << ir::decl_to_string(d) << "\n";
     tout() << "Max var: " << ir::decl_max_var(d).to_std_string() << "\n";
+    tout() << "after push proj\n";
+    d = ir::decl_push_proj(d);
+    tout() << ir::decl_to_string(d) << "\n";
 }
 
 pair<environment, comp_decls> to_llnf(environment const & env, comp_decls const & ds, bool unboxed) {