From cdf79d07785bc84f001b732daadcadecbe7527bd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?William=20S=C3=B8rensen?= <ws423@cam.ac.uk>
Date: Wed, 26 Jun 2024 11:10:19 +0100
Subject: [PATCH 1/4] refactor: move elabQpf into seperate funs

---
 Qpf/Macro/Comp.lean | 165 ++++++++++++++++++++++++--------------------
 1 file changed, 91 insertions(+), 74 deletions(-)

diff --git a/Qpf/Macro/Comp.lean b/Qpf/Macro/Comp.lean
index 65ce150..8d10d00 100644
--- a/Qpf/Macro/Comp.lean
+++ b/Qpf/Macro/Comp.lean
@@ -160,97 +160,115 @@ structure ElabQpfResult (u : Level) (arity : Nat) where
 deriving Inhabited
 
 open PrettyPrinter in
-/--
-  Elaborate the body of a qpf
--/
-partial def elabQpf {arity : Nat} (vars : Vector Q(Type u) arity) (target : Q(Type u)) (targetStx : Option Term := none) (normalized := false) :
-    TermElabM (ElabQpfResult u arity) := do
-  trace[QPF] "elabQPF :: {vars.toList} -> {target}"
+mutual
+partial def handleLiveFVar (vars : Vector Q(Type u) arity) (target : Q(Type u))  : TermElabM (ElabQpfResult u arity) := do
+  let vars' := vars.toList;
+
+  trace[QPF] f!"target {target} is a free variable"
+  let ind ← match List.indexOf' target vars' with
+  | none      => throwError "Free variable {target} is not one of the qpf arguments"
+  | some ind  => pure ind
+
+  let ind : Fin2 arity := cast (by simp [vars']) ind.inv
+  let prj := q(@Prj.{u} $arity $ind)
+  trace[QPF] "represented by: {prj}"
+  pure ⟨prj, q(Prj.mvfunctor _), q(Prj.mvqpf _)⟩
+
+partial def handleConst (vars : Vector Q(Type u) arity) (target : Q(Type u))  : TermElabM (ElabQpfResult u arity) := do
+  trace[QPF] "target {target} is a constant"
+  let const := q(Const.{u} $arity $target)
+  trace[QPF] "represented by: {const}"
+  pure ⟨const, q(Const.MvFunctor), q(Const.mvqpf)⟩
+
+partial def handleApp (vars : Vector Q(Type u) arity) (target : Q(Type u))  : TermElabM (ElabQpfResult u arity) := do
   let vars' := vars.toList;
 
   let varIds := vars'.map fun expr => expr.fvarId!
   let isLiveVar : FVarId → Bool
     := fun fvarId => (List.indexOf' fvarId varIds).isSome
 
-  if target.isFVar && isLiveVar target.fvarId! then
-    trace[QPF] f!"target {target} is a free variable"
-    let ind ← match List.indexOf' target vars' with
-    | none      => throwError "Free variable {target} is not one of the qpf arguments"
-    | some ind  => pure ind
-
-    let ind : Fin2 arity := cast (by simp [vars']) ind.inv
-    let prj := q(@Prj.{u} $arity $ind)
-    trace[QPF] "represented by: {prj}"
-    pure ⟨prj, q(Prj.mvfunctor _), q(Prj.mvqpf _)⟩
-
-  else if !target.hasAnyFVar isLiveVar then
-    trace[QPF] "target {target} is a constant"
-    let const := q(Const.{u} $arity $target)
-    trace[QPF] "represented by: {const}"
-    pure ⟨const, q(Const.MvFunctor), q(Const.mvqpf)⟩
+  let ⟨m, F, args⟩ ← (Comp.parseApp isLiveVar target)
+  trace[QPF] "target {target} is an application of {F} to {args.toList}"
 
-  else if target.isApp then
-    let ⟨m, F, args⟩ ← (Comp.parseApp isLiveVar target)
-    trace[QPF] "target {target} is an application of {F} to {args.toList}"
+  /-
+    Optimization: check if the application is of the form `F α β γ .. = F α β γ ..`.
+    In such cases, we can directly return `F`, rather than generate a composition of projections.
+  -/
+  let is_trivial :=
+    args.length == arity
+    && args.toList.enum.all fun ⟨i, arg⟩ =>
+        arg.isFVar && isLiveVar arg.fvarId! && vars'.indexOf arg == i
+  if is_trivial then
+    trace[QPF] "The application is trivial"
+    let mvFunctor ← synthInstanceQ q(MvFunctor $F)
+    let mvQPF ← synthInstanceQ q(MvQPF $F)
+    pure ⟨F, mvFunctor, mvQPF⟩
+  else
+    let G ← args.toList.mapM fun a =>
+      elabQpf vars a none false
 
     /-
-      Optimization: check if the application is of the form `F α β γ .. = F α β γ ..`.
-      In such cases, we can directly return `F`, rather than generate a composition of projections.
+      HACK: We know that `m`, which was defines as `args.length`, is equal to `G.length`.
+      It's a bit difficult to prove this. Thus, we simply assert it
     -/
-    let is_trivial :=
-      args.length == arity
-      && args.toList.enum.all fun ⟨i, arg⟩ =>
-          arg.isFVar && isLiveVar arg.fvarId! && vars'.indexOf arg == i
-    if is_trivial then
-      trace[QPF] "The application is trivial"
-      let mvFunctor ← synthInstanceQ q(MvFunctor $F)
-      let mvQPF ← synthInstanceQ q(MvQPF $F)
-      pure ⟨F, mvFunctor, mvQPF⟩
+    if hm : m ≠ G.length then
+      throwError "This shouldn't happen" -- TODO: come up with a better error message
     else
-      let G ← args.toList.mapM fun a =>
-        elabQpf vars a none false
-
-      /-
-        HACK: We know that `m`, which was defines as `args.length`, is equal to `G.length`.
-        It's a bit difficult to prove this. Thus, we simply assert it
-      -/
-      if hm : m ≠ G.length then
-        throwError "This shouldn't happen" -- TODO: come up with a better error message
-      else
-        have hm : m = G.length := by simpa using hm
+      have hm : m = G.length := by simpa using hm
 
 
 
-        let Ffunctor ← synthInstanceQ q(MvFunctor $F)
-        let Fqpf ← synthInstanceQ q(@MvQPF _ $F $Ffunctor)
+      let Ffunctor ← synthInstanceQ q(MvFunctor $F)
+      let Fqpf ← synthInstanceQ q(@MvQPF _ $F $Ffunctor)
 
-        let G : Vec _ m := fun i => G.get (hm ▸ i.inv)
-        let GExpr : Q(Vec (TypeFun.{u,u} $arity) $m) :=
-          Vec.toExpr (fun i => (G i).F)
-        let GFunctor : Q(∀ i, MvFunctor.{u,u} ($GExpr i)) :=
-          let αs := q(fun i => MvFunctor.{u,u} ($GExpr i))
-          @DVec.toExpr _ _ αs (fun i => (G i).functor)
-        let GQpf : Q(∀ i, @MvQPF.{u,u} _ _ ($GFunctor i)) :=
-          let αs := q(fun i => @MvQPF.{u,u} _ _ ($GFunctor i))
-          @DVec.toExpr _ _ αs (fun i => (G i).qpf)
+      let G : Vec _ m := fun i => G.get (hm ▸ i.inv)
+      let GExpr : Q(Vec (TypeFun.{u,u} $arity) $m) :=
+        Vec.toExpr (fun i => (G i).F)
+      let GFunctor : Q(∀ i, MvFunctor.{u,u} ($GExpr i)) :=
+        let αs := q(fun i => MvFunctor.{u,u} ($GExpr i))
+        @DVec.toExpr _ _ αs (fun i => (G i).functor)
+      let GQpf : Q(∀ i, @MvQPF.{u,u} _ _ ($GFunctor i)) :=
+        let αs := q(fun i => @MvQPF.{u,u} _ _ ($GFunctor i))
+        @DVec.toExpr _ _ αs (fun i => (G i).qpf)
 
-        let comp := q(@Comp $m _ $F $GExpr)
-        trace[QPF] "G := {GExpr}"
-        trace[QPF] "comp := {comp}"
+      let comp := q(@Comp $m _ $F $GExpr)
+      trace[QPF] "G := {GExpr}"
+      trace[QPF] "comp := {comp}"
 
-        let functor := q(Comp.instMvFunctorComp)
-        let qpf := q(Comp.instMvQPFCompInstMvFunctorCompFin2
-          (fF := $Ffunctor) (q := $Fqpf) (fG := _) (q' := $GQpf)
-        )
-        pure ⟨comp, functor, qpf⟩
+      let functor := q(Comp.instMvFunctorComp)
+      let qpf := q(Comp.instMvQPFCompInstMvFunctorCompFin2
+        (fF := $Ffunctor) (q := $Fqpf) (fG := _) (q' := $GQpf)
+      )
+      pure ⟨comp, functor, qpf⟩
 
-  else if target.isArrow then
-    match target with
-    | Expr.forallE _ e₁ e₂ .. =>
-      let newTarget ← mkAppM ``MvQPF.Arrow.Arrow #[e₁, e₂]
-      elabQpf vars newTarget targetStx normalized
-    | _ => unreachable!
 
+partial def handleArrow (vars : Vector Q(Type u) arity) (target : Q(Type u)) (targetStx : Option Term := none) (normalized := false) : TermElabM (ElabQpfResult u arity) := do 
+  match target with
+  | Expr.forallE _ e₁ e₂ .. =>
+    let newTarget ← mkAppM ``MvQPF.Arrow.Arrow #[e₁, e₂]
+    elabQpf vars newTarget targetStx normalized
+  | _ => unreachable!
+
+/--
+  Elaborate the body of a qpf
+-/
+partial def elabQpf {arity : Nat} (vars : Vector Q(Type u) arity) (target : Q(Type u)) (targetStx : Option Term := none) (normalized := false) :
+    TermElabM (ElabQpfResult u arity) := do
+  trace[QPF] "elabQPF :: {vars.toList} -> {target}"
+  let vars' := vars.toList;
+
+  let varIds := vars'.map fun expr => expr.fvarId!
+  let isLiveVar : FVarId → Bool
+    := fun fvarId => (List.indexOf' fvarId varIds).isSome
+
+  if target.isFVar && isLiveVar target.fvarId! then
+    handleLiveFVar vars target
+  else if !target.hasAnyFVar isLiveVar then
+    handleConst vars target
+  else if target.isApp then
+    handleApp vars target
+  else if target.isArrow then
+    handleArrow vars target (targetStx := targetStx) (normalized := normalized)
   else
     if !normalized then
       let target ← whnfR target
@@ -262,8 +280,7 @@ partial def elabQpf {arity : Nat} (vars : Vector Q(Type u) arity) (target : Q(Ty
         else
           ""
       throwError f!"Unexpected target expression :\n {target}\n{extra}\nNote that the expression contains live variables, hence, must be functorial"
-
-
+end
 
 
 structure QpfCompositionBodyView where

From 1720c1efdc9850cc271e0d97ad18805e4245cd79 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?William=20S=C3=B8rensen?= <ws423@cam.ac.uk>
Date: Wed, 26 Jun 2024 11:20:46 +0100
Subject: [PATCH 2/4] refactor: update types in elabQpf to reflect the code

---
 Qpf/Macro/Comp.lean | 175 +++++++++++++++++++++-----------------------
 1 file changed, 85 insertions(+), 90 deletions(-)

diff --git a/Qpf/Macro/Comp.lean b/Qpf/Macro/Comp.lean
index 8d10d00..1d8faef 100644
--- a/Qpf/Macro/Comp.lean
+++ b/Qpf/Macro/Comp.lean
@@ -159,130 +159,121 @@ structure ElabQpfResult (u : Level) (arity : Nat) where
   qpf     : Q(@MvQPF _ $F $functor)
 deriving Inhabited
 
+-- TODO: Move to Vector.mmap
+def Vector.mapM {m : Type u → Type v}
+    {α : Type w} {β : Type u} [Monad m]
+    (l : Vector α n) (f : α → m β) : m $ Vector β n :=
+  match l with
+  | ⟨.cons hd tl, b⟩ => do
+    let v ← f hd
+    let x ← Vector.mapM ⟨tl, rfl⟩ f
+    pure ⟨v :: x.val, by simpa [List.length_cons, x.property] using b⟩
+  | ⟨.nil, h⟩ => pure ⟨[], by exact h⟩
+
+def isLiveVar (varIds : Vector FVarId n) (id : FVarId) := varIds.toList.contains id
+
 open PrettyPrinter in
 mutual
-partial def handleLiveFVar (vars : Vector Q(Type u) arity) (target : Q(Type u))  : TermElabM (ElabQpfResult u arity) := do
-  let vars' := vars.toList;
-
-  trace[QPF] f!"target {target} is a free variable"
-  let ind ← match List.indexOf' target vars' with
-  | none      => throwError "Free variable {target} is not one of the qpf arguments"
+partial def handleLiveFVar (vars : Vector FVarId arity) (target : FVarId)  : TermElabM (ElabQpfResult u arity) := do
+  trace[QPF] f!"target {Expr.fvar target} is a free variable"
+  let ind ← match List.indexOf' target vars.toList with
+  | none      => throwError "Free variable {Expr.fvar target} is not one of the qpf arguments"
   | some ind  => pure ind
 
-  let ind : Fin2 arity := cast (by simp [vars']) ind.inv
+  let ind : Fin2 arity := cast (by simp) ind.inv
   let prj := q(@Prj.{u} $arity $ind)
   trace[QPF] "represented by: {prj}"
-  pure ⟨prj, q(Prj.mvfunctor _), q(Prj.mvqpf _)⟩
 
-partial def handleConst (vars : Vector Q(Type u) arity) (target : Q(Type u))  : TermElabM (ElabQpfResult u arity) := do
+  pure { F := prj, functor := q(Prj.mvfunctor _), qpf := q(Prj.mvqpf _) }
+
+partial def handleConst (target : Q(Type u))  : TermElabM (ElabQpfResult u arity) := do
   trace[QPF] "target {target} is a constant"
   let const := q(Const.{u} $arity $target)
   trace[QPF] "represented by: {const}"
-  pure ⟨const, q(Const.MvFunctor), q(Const.mvqpf)⟩
 
-partial def handleApp (vars : Vector Q(Type u) arity) (target : Q(Type u))  : TermElabM (ElabQpfResult u arity) := do
-  let vars' := vars.toList;
+  pure { F := const, functor := q(Const.MvFunctor), qpf := q(Const.mvqpf)}
 
-  let varIds := vars'.map fun expr => expr.fvarId!
-  let isLiveVar : FVarId → Bool
-    := fun fvarId => (List.indexOf' fvarId varIds).isSome
+partial def handleApp (vars : Vector FVarId arity) (target : Q(Type u))  : TermElabM (ElabQpfResult u arity) := do
+  let vars' := vars.toList
 
-  let ⟨m, F, args⟩ ← (Comp.parseApp isLiveVar target)
+  let ⟨numArgs, F, args⟩ ← (Comp.parseApp (isLiveVar vars) target)
   trace[QPF] "target {target} is an application of {F} to {args.toList}"
 
   /-
-    Optimization: check if the application is of the form `F α β γ .. = F α β γ ..`.
-    In such cases, we can directly return `F`, rather than generate a composition of projections.
+    Optimization: check if the application is of the form `F α β γ .. = G α β γ ..`.
+    In such cases, we can directly return `G`, rather than generate a composition of projections.
   -/
-  let is_trivial :=
-    args.length == arity
-    && args.toList.enum.all fun ⟨i, arg⟩ =>
-        arg.isFVar && isLiveVar arg.fvarId! && vars'.indexOf arg == i
-  if is_trivial then
-    trace[QPF] "The application is trivial"
-    let mvFunctor ← synthInstanceQ q(MvFunctor $F)
-    let mvQPF ← synthInstanceQ q(MvQPF $F)
-    pure ⟨F, mvFunctor, mvQPF⟩
-  else
-    let G ← args.toList.mapM fun a =>
-      elabQpf vars a none false
-
-    /-
-      HACK: We know that `m`, which was defines as `args.length`, is equal to `G.length`.
-      It's a bit difficult to prove this. Thus, we simply assert it
-    -/
-    if hm : m ≠ G.length then
-      throwError "This shouldn't happen" -- TODO: come up with a better error message
-    else
-      have hm : m = G.length := by simpa using hm
-
 
+  -- O(n²), equivalent to a zipping and would be O(n) and much more readable
+  /- let is_trivial := -/
+  /-   args.length == arity -/
+  /-   && args.toList.enum.all fun ⟨i, arg⟩ => -/
+  -- Equivalent (?), O(n), and more readable
+  -- TODO: is this really equivalent
+  let is_trivial := (args.toList.mapM Expr.fvarId?).any (· == vars')
+  /- let is_trivial := (args.toList.mapM Expr.fvarId?).any (BEq.beq vars') -/
 
-      let Ffunctor ← synthInstanceQ q(MvFunctor $F)
-      let Fqpf ← synthInstanceQ q(@MvQPF _ $F $Ffunctor)
-
-      let G : Vec _ m := fun i => G.get (hm ▸ i.inv)
-      let GExpr : Q(Vec (TypeFun.{u,u} $arity) $m) :=
-        Vec.toExpr (fun i => (G i).F)
-      let GFunctor : Q(∀ i, MvFunctor.{u,u} ($GExpr i)) :=
-        let αs := q(fun i => MvFunctor.{u,u} ($GExpr i))
-        @DVec.toExpr _ _ αs (fun i => (G i).functor)
-      let GQpf : Q(∀ i, @MvQPF.{u,u} _ _ ($GFunctor i)) :=
-        let αs := q(fun i => @MvQPF.{u,u} _ _ ($GFunctor i))
-        @DVec.toExpr _ _ αs (fun i => (G i).qpf)
-
-      let comp := q(@Comp $m _ $F $GExpr)
-      trace[QPF] "G := {GExpr}"
-      trace[QPF] "comp := {comp}"
+  if is_trivial then
+    trace[QPF] "The application is trivial"
+    let functor ← synthInstanceQ q(MvFunctor $F)
+    let qpf ← synthInstanceQ q(MvQPF $F)
 
-      let functor := q(Comp.instMvFunctorComp)
-      let qpf := q(Comp.instMvQPFCompInstMvFunctorCompFin2
-        (fF := $Ffunctor) (q := $Fqpf) (fG := _) (q' := $GQpf)
-      )
-      pure ⟨comp, functor, qpf⟩
+    return { F, functor, qpf }
+  else
+    let G ← Vector.mapM args (elabQpf vars · none false)
+
+    let Ffunctor ← synthInstanceQ q(MvFunctor $F)
+    let Fqpf ← synthInstanceQ q(@MvQPF _ $F $Ffunctor)
+
+    let G : Vec _ numArgs := fun i => G.get i.inv
+    let GExpr : Q(Vec (TypeFun.{u,u} $arity) $numArgs) :=
+      Vec.toExpr (fun i => (G i).F)
+    let GFunctor : Q(∀ i, MvFunctor.{u,u} ($GExpr i)) :=
+      let αs := q(fun i => MvFunctor.{u,u} ($GExpr i))
+      @DVec.toExpr _ _ αs (fun i => (G i).functor)
+    let GQpf : Q(∀ i, @MvQPF.{u,u} _ _ ($GFunctor i)) :=
+      let αs := q(fun i => @MvQPF.{u,u} _ _ ($GFunctor i))
+      @DVec.toExpr _ _ αs (fun i => (G i).qpf)
+
+    let comp := q(@Comp $numArgs _ $F $GExpr)
+    trace[QPF] "G := {GExpr}"
+    trace[QPF] "comp := {comp}"
+
+    let functor := q(Comp.instMvFunctorComp)
+    let qpf := q(Comp.instMvQPFCompInstMvFunctorCompFin2
+      (fF := $Ffunctor) (q := $Fqpf) (fG := _) (q' := $GQpf)
+    )
 
+    return { F := comp, functor, qpf }
 
-partial def handleArrow (vars : Vector Q(Type u) arity) (target : Q(Type u)) (targetStx : Option Term := none) (normalized := false) : TermElabM (ElabQpfResult u arity) := do 
-  match target with
-  | Expr.forallE _ e₁ e₂ .. =>
-    let newTarget ← mkAppM ``MvQPF.Arrow.Arrow #[e₁, e₂]
-    elabQpf vars newTarget targetStx normalized
-  | _ => unreachable!
+partial def handleArrow (binderType body : Expr) (vars : Vector FVarId arity) (targetStx : Option Term := none) (normalized := false): TermElabM (ElabQpfResult u arity) := do
+  let newTarget ← mkAppM ``MvQPF.Arrow.Arrow #[binderType, body]
+  elabQpf vars newTarget targetStx normalized
 
 /--
   Elaborate the body of a qpf
 -/
-partial def elabQpf {arity : Nat} (vars : Vector Q(Type u) arity) (target : Q(Type u)) (targetStx : Option Term := none) (normalized := false) :
+partial def elabQpf {arity : Nat} (vars : Vector FVarId arity) (target : Q(Type u)) (targetStx : Option Term := none) (normalized := false) :
     TermElabM (ElabQpfResult u arity) := do
-  trace[QPF] "elabQPF :: {vars.toList} -> {target}"
-  let vars' := vars.toList;
+  trace[QPF] "elabQPF :: {(vars.map Expr.fvar).toList} -> {target}"
+  let isLiveVar := isLiveVar vars
 
-  let varIds := vars'.map fun expr => expr.fvarId!
-  let isLiveVar : FVarId → Bool
-    := fun fvarId => (List.indexOf' fvarId varIds).isSome
-
-  if target.isFVar && isLiveVar target.fvarId! then
+  if let some target := target.fvarId?.filter isLiveVar then
     handleLiveFVar vars target
   else if !target.hasAnyFVar isLiveVar then
-    handleConst vars target
-  else if target.isApp then
+    handleConst target
+  else if target.isApp then -- Could be pattern-matched here as well
     handleApp vars target
-  else if target.isArrow then
-    handleArrow vars target (targetStx := targetStx) (normalized := normalized)
+  else if let .forallE _ binderType body .. := target then
+    handleArrow binderType body vars (normalized := normalized) (targetStx := targetStx)
+  else if !normalized then
+    let target ← whnfR target
+    elabQpf vars target targetStx true
   else
-    if !normalized then
-      let target ← whnfR target
-      elabQpf vars target targetStx true
-    else
-      let extra :=
-        if target.isForall then
-          "Dependent arrows / forall are not supported"
-        else
-          ""
-      throwError f!"Unexpected target expression :\n {target}\n{extra}\nNote that the expression contains live variables, hence, must be functorial"
+    let extra := if target.isForall then "Dependent arrows / forall are not supported" else ""
+    throwError f!"Unexpected target expression :\n {target}\n{extra}\nNote that the expression contains live variables, hence, must be functorial"
 end
 
-
 structure QpfCompositionBodyView where
   (type?  : Option Syntax := none)
   (target : Term)
@@ -323,6 +314,10 @@ def elabQpfCompositionBody (view: QpfCompositionBodyView) :
             let target_expr ← elabTermEnsuringTypeQ (u:=u.succ.succ) view.target q(Type u)
             let arity := vars.toList.length
             let vars : Vector _ arity := ⟨vars.toList, rfl⟩
+
+            let some vars := Vector.mapM vars Expr.fvarId? |
+              throwError "Expected all args to be fvars"
+
             let res ← elabQpf vars target_expr view.target
 
             res.F.check

From bfc07e0a3e9b64f14157e1e76fe0dae9c7ed8533 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?William=20S=C3=B8rensen?= <ws423@cam.ac.uk>
Date: Wed, 26 Jun 2024 13:53:05 +0100
Subject: [PATCH 3/4] refactor: revert to old is_trivial

---
 Qpf/Macro/Comp.lean | 34 ++++++++++++----------------------
 1 file changed, 12 insertions(+), 22 deletions(-)

diff --git a/Qpf/Macro/Comp.lean b/Qpf/Macro/Comp.lean
index 322bd8a..a7336e5 100644
--- a/Qpf/Macro/Comp.lean
+++ b/Qpf/Macro/Comp.lean
@@ -154,22 +154,11 @@ def DVec.toExpr {n : Nat} {αs : Q(Fin2 $n → Type u)} (xs : DVec (fun (i : Fin
 
 
 structure ElabQpfResult (u : Level) (arity : Nat) where
-  F : Q(TypeFun.{u,u} $arity)
-  functor : Q(MvFunctor $F)
-  qpf     : Q(@MvQPF _ $F $functor)
+  F : Q(TypeFun.{u,u} $arity)       := by exact q(by infer_instance)
+  functor : Q(MvFunctor $F)         := by exact q(by infer_instance)
+  qpf     : Q(@MvQPF _ $F $functor) := by exact q(by infer_instance)
 deriving Inhabited
 
--- TODO: Move to Vector.mmap
-def Vector.mapM {m : Type u → Type v}
-    {α : Type w} {β : Type u} [Monad m]
-    (l : Vector α n) (f : α → m β) : m $ Vector β n :=
-  match l with
-  | ⟨.cons hd tl, b⟩ => do
-    let v ← f hd
-    let x ← Vector.mapM ⟨tl, rfl⟩ f
-    pure ⟨v :: x.val, by simpa [List.length_cons, x.property] using b⟩
-  | ⟨.nil, h⟩ => pure ⟨[], by exact h⟩
-
 def isLiveVar (varIds : Vector FVarId n) (id : FVarId) := varIds.toList.contains id
 
 open PrettyPrinter in
@@ -193,6 +182,7 @@ partial def handleConst (target : Q(Type u))  : TermElabM (ElabQpfResult u arity
 
   pure { F := const, functor := q(Const.MvFunctor), qpf := q(Const.mvqpf)}
 
+
 partial def handleApp (vars : Vector FVarId arity) (target : Q(Type u))  : TermElabM (ElabQpfResult u arity) := do
   let vars' := vars.toList
 
@@ -205,13 +195,13 @@ partial def handleApp (vars : Vector FVarId arity) (target : Q(Type u))  : TermE
   -/
 
   -- O(n²), equivalent to a zipping and would be O(n) and much more readable
-  /- let is_trivial := -/
-  /-   args.length == arity -/
-  /-   && args.toList.enum.all fun ⟨i, arg⟩ => -/
+  let is_trivial :=
+    args.length == arity
+    && args.toList.enum.all fun ⟨i, arg⟩ =>
+        arg.isFVar && isLiveVar vars arg.fvarId! && vars'.indexOf arg.fvarId! == i
   -- Equivalent (?), O(n), and more readable
   -- TODO: is this really equivalent
-  let is_trivial := (args.toList.mapM Expr.fvarId?).any (· == vars')
-  /- let is_trivial := (args.toList.mapM Expr.fvarId?).any (BEq.beq vars') -/
+  /- let is_trivial := (args.toList.mapM Expr.fvarId?).any (· == vars') -/
 
   if is_trivial then
     trace[QPF] "The application is trivial"
@@ -220,7 +210,7 @@ partial def handleApp (vars : Vector FVarId arity) (target : Q(Type u))  : TermE
 
     return { F, functor, qpf }
   else
-    let G ← Vector.mapM args (elabQpf vars · none false)
+    let G ← args.mmap (elabQpf vars · none false)
 
     let Ffunctor ← synthInstanceQ q(MvFunctor $F)
     let Fqpf ← synthInstanceQ q(@MvQPF _ $F $Ffunctor)
@@ -314,8 +304,8 @@ def elabQpfCompositionBody (view: QpfCompositionBodyView) :
             let arity := vars.toList.length
             let vars : Vector _ arity := ⟨vars.toList, rfl⟩
 
-            let some vars := Vector.mapM vars Expr.fvarId? |
-              throwError "Expected all args to be fvars"
+            let some vars := vars.mmap Expr.fvarId? |
+              throwError "Expected all args to be fvars" -- TODO: throwErrorAt
 
             let res ← elabQpf vars target_expr view.target
 

From 76af185c57b88b3c0af43c2ea092a32a1cec3367 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?William=20S=C3=B8rensen?=
 <47296141+Equilibris@users.noreply.github.com>
Date: Fri, 28 Jun 2024 17:48:39 +0200
Subject: [PATCH 4/4] Apply suggestions from code review

Fix whitespace
---
 Qpf/Macro/Comp.lean | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/Qpf/Macro/Comp.lean b/Qpf/Macro/Comp.lean
index 68aa414..35a2a4a 100644
--- a/Qpf/Macro/Comp.lean
+++ b/Qpf/Macro/Comp.lean
@@ -154,7 +154,7 @@ def DVec.toExpr {n : Nat} {αs : Q(Fin2 $n → Type u)} (xs : DVec (fun (i : Fin
 
 
 structure ElabQpfResult (u : Level) (arity : Nat) where
-  F : Q(TypeFun.{u,u} $arity)       := by exact q(by infer_instance)
+  F : Q(TypeFun.{u,u} $arity)
   functor : Q(MvFunctor $F)         := by exact q(by infer_instance)
   qpf     : Q(@MvQPF _ $F $functor) := by exact q(by infer_instance)
 deriving Inhabited
@@ -183,7 +183,7 @@ partial def handleConst (target : Q(Type u))  : TermElabM (ElabQpfResult u arity
   pure { F := const, functor := q(Const.MvFunctor), qpf := q(Const.mvqpf)}
 
 
-partial def handleApp (vars : Vector FVarId arity) (target : Q(Type u))  : TermElabM (ElabQpfResult u arity) := do
+partial def handleApp (vars : Vector FVarId arity) (target : Q(Type u)) : TermElabM (ElabQpfResult u arity) := do
   let vars' := vars.toList
 
   let ⟨numArgs, F, args⟩ ← (Comp.parseApp (isLiveVar vars) target)
@@ -236,7 +236,7 @@ partial def handleApp (vars : Vector FVarId arity) (target : Q(Type u))  : TermE
     return { F := comp, functor, qpf }
 
 
-partial def handleArrow (binderType body : Expr) (vars : Vector FVarId arity) (targetStx : Option Term := none) (normalized := false): TermElabM (ElabQpfResult u arity) := do
+partial def handleArrow (binderType body : Expr) (vars : Vector FVarId arity) (targetStx : Option Term := none) (normalized := false) : TermElabM (ElabQpfResult u arity) := do
   let newTarget ← mkAppM ``MvQPF.Arrow.Arrow #[binderType, body]
   elabQpf vars newTarget targetStx normalized
 /--