Merge branch 'master' into refactor-replace

alexkeizer · Jun 26, 2024 · 2c17e2a · 2c17e2a
2 parents 33b6353 + c62f9cc
commit 2c17e2a
Showing 1 changed file with 91 additions and 74 deletions.
diff --git 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