From a36c4383baac503eb8d836837955878846bf2a97 Mon Sep 17 00:00:00 2001
From: lecopivo <skrivantomas@seznam.cz>
Date: Thu, 21 Sep 2023 16:22:05 -0400
Subject: [PATCH] meta functions to extend local context to make sure types
 have instances of Vec, SemiInnerProductSpace etc.

---
 SciLean/Core/Meta/ExtendContext.lean     | 230 +++++++++++++++++++++++
 SciLean/Core/Meta/GenerateBasic.lean     |  42 +++++
 SciLean/Core/Meta/GenerateRevCDeriv.lean |  49 +----
 test/extend_context.lean                 |  46 +++++
 4 files changed, 322 insertions(+), 45 deletions(-)
 create mode 100644 SciLean/Core/Meta/ExtendContext.lean
 create mode 100644 test/extend_context.lean

diff --git a/SciLean/Core/Meta/ExtendContext.lean b/SciLean/Core/Meta/ExtendContext.lean
new file mode 100644
index 00000000..e2c6e081
--- /dev/null
+++ b/SciLean/Core/Meta/ExtendContext.lean
@@ -0,0 +1,230 @@
+import SciLean.Core.Objects.FinVec
+import SciLean.Lean.Meta.Basic
+import SciLean.Data.Index
+
+open Lean Meta Qq
+
+namespace SciLean
+
+
+/-- Modifies the local context such that `I` has instance `EnumType I`
+
+All newly introduced free variables are passed to k. 
+
+If necessary it introduces
+- `EnumType J` if `J` is fvar
+-/
+partial def withEnumType {α : Type _}
+  (I : Expr) (v w : Level) (k : Array Expr → MetaM α) : MetaM α := do
+  loop I #[] k
+where
+  loop (I : Expr) (acc : Array Expr) (k : Array Expr → MetaM α) : MetaM α := do
+    let .some ⟨u,_⟩ ← isTypeQ I | throwError "invalid type {← ppExpr I}"
+    let cls := (Expr.const ``EnumType [u,v,w]).app I
+    match ← synthInstance? cls with
+    | .some _ => k acc
+    | none => 
+      match I with
+      | .app .. => 
+        if (I.isAppOfArity' ``Prod 2) ||
+           (I.isAppOfArity' ``ColProd 2) ||
+           (I.isAppOfArity' ``Sum 2) then
+          let X₁ := I.getArg! 0
+          let X₂ := I.getArg! 1
+          loop X₁ acc (fun acc' => loop X₂ acc' k)
+        else
+          throwError "dont' know how to extend context to have instance `{← ppExpr cls}`"
+      | .fvar _ => 
+        withLocalDecl (← mkAuxName `inst 0) .instImplicit cls fun inst => 
+          k (acc.push inst)
+      | _ => 
+        throwError "dont' know how to extend context to have instance `{← ppExpr cls}`"
+
+
+/-- Modifies the local context such that `I` has instance `Index I`
+
+All newly introduced free variables are passed to k. 
+
+If necessary it introduces
+- `Index J` if `J` is fvar
+
+It modifies existing instances 
+- `EnumType J` to `Index J`
+-/
+partial def withIndex {α : Type _}
+  (I : Expr) (v w : Level) (k : Array Expr → MetaM α) : MetaM α := do
+  loop I #[] k
+where
+  loop (I : Expr) (acc : Array Expr) (k : Array Expr → MetaM α) : MetaM α := do
+    let .some ⟨u,_⟩ ← isTypeQ I | throwError "invalid type {← ppExpr I}"
+    let cls := (Expr.const ``Index [u,v,w]).app I
+    match ← synthInstance? cls with
+    | .some _ => k acc
+    | none => 
+      match I with
+      | .app .. => 
+        if (I.isAppOfArity' ``Prod 2) ||
+           (I.isAppOfArity' ``ColProd 2) ||
+           (I.isAppOfArity' ``Sum 2) then
+          let X₁ := I.getArg! 0
+          let X₂ := I.getArg! 1
+          loop X₁ acc (fun acc' => loop X₂ acc' k)
+        else
+          throwError "dont' know how to extend context to have instance `{← ppExpr cls}`"
+      | .fvar _ => 
+        withLocalDecl (← mkAuxName `inst 0) .instImplicit cls fun inst => 
+          k (acc.push inst)
+      | _ => 
+        throwError "dont' know how to extend context to have instance `{← ppExpr cls}`"
+
+
+/-- Modifies the local context such that `X` has instance `Vec K X`
+
+All newly introduced free variables are passed to k. 
+
+If necessary it introduces
+- `Vec K X` if `X` is fvar
+-/
+partial def withVec {α : Type _}
+  (K X : Expr) (k : Array Expr → MetaM α) : MetaM α := do
+  let .some ⟨_u,K⟩ ← isTypeQ K | throwError "invalid type {← ppExpr K}"
+  loop K X #[] k
+where
+  loop {u} (K : Q(Type $u)) (X : Expr) (acc : Array Expr) (k : Array Expr → MetaM α) : MetaM α := do
+    let cls ← mkAppM ``Vec #[K, X]
+    match ← synthInstance? cls with
+    | .some _ => k acc
+    | none => 
+      match X with
+      | .forallE _ _ Y _ => 
+        if Y.hasLooseBVars then
+          throwError "can't extend context for dependent type `{← ppExpr X}`"
+        loop K Y acc k
+      | .app .. => 
+        if X.isAppOfArity' ``Prod 2 then
+          let X₁ := X.getArg! 0
+          let X₂ := X.getArg! 1
+          loop K X₁ acc (fun acc' => loop K X₂ acc' k)
+        else
+          throwError "dont' know how to extend context for the type `{← ppExpr X}`"
+      | .fvar _ => 
+        withLocalDecl (← mkAuxName `inst 0) .instImplicit cls fun inst => 
+          k (acc.push inst)
+      | _ => 
+        throwError "dont' know how to extend context for the type `{← ppExpr X}`"
+
+
+/-- Modifies the local context such that `X` has instance `SemiInnerProductSpace K X`
+
+All newly introduced free variables are passed to k. 
+
+If necessary it introduces
+- `EnumType I` for `X = I → Y` 
+- `SemiInnerProductSpace K X` if `X` is fvar
+
+It modifies existing instances 
+- `Vec K X` to `SemiInnerProductSpace K X`
+-/
+partial def withSemiInnerProductSpace {α : Type _}
+  (K X : Expr) (k : Array Expr → MetaM α) : MetaM α := do
+  let .some ⟨_u,K⟩ ← isTypeQ K | throwError "invalid type {← ppExpr K}"
+  loop K X #[] k
+where
+  loop {u} (K : Q(Type $u)) (X : Expr) (acc : Array Expr) (k : Array Expr → MetaM α) : MetaM α := do
+    let cls ← mkAppM ``SemiInnerProductSpace #[K, X]
+    match ← synthInstance? cls with
+    | .some _ => k acc
+    | none => 
+      match X with
+      | .forallE _ I Y _ => 
+        if Y.hasLooseBVars then
+          throwError "can't extend context for dependent type `{← ppExpr X}`"
+        withEnumType I u u (fun acc' => loop K Y (acc ++ acc') k)
+      | .app .. => 
+        if X.isAppOfArity' ``Prod 2 then
+          let X₁ := X.getArg! 0
+          let X₂ := X.getArg! 1
+          loop K X₁ acc (fun acc' => loop K X₂ acc' k)
+        else
+          throwError "dont' know how to extend context for the type `{← ppExpr X}`"
+      | .fvar _ => 
+        -- try to upgrade Vec to SemiInnerProductSpace
+        let vecX ← mkAppM ``Vec #[K,X]
+        let lctx ← getLCtx
+        let vecId? ← lctx.findDeclM? 
+          (fun decl => do 
+            pure <| if (← isDefEq decl.type vecX) then .some (decl.fvarId) else .none)
+        match vecId? with
+        | .some vecId => 
+            let lctx' := lctx.modifyLocalDecl vecId
+              fun decl => decl.setType cls
+            let insts ← getLocalInstances
+            let insts' := insts.erase vecId 
+              |>.push {className := ``SemiInnerProductSpace, fvar := .fvar vecId}
+            withLCtx lctx' insts' (k acc)
+        | .none => 
+          withLocalDecl (← mkAuxName `inst 0) .instImplicit cls fun inst => 
+            k (acc.push inst)
+      | _ => 
+        throwError "dont' know how to extend context for the type `{← ppExpr X}`"
+
+
+/-- Modifies the local context such that `X` has instance `SemiHilbert K X`
+
+All newly introduced free variables are passed to k. 
+
+If necessary it introduces
+- `EnumType I` for `X = I → Y` 
+- `SemiHilbert K X` if `X` is fvar
+
+It modifies
+- instance `Vec K X` to `SemiHilbert K X`
+- instance `SemiInnerProductSpace K X` to `SemiHilbert K X`
+-/
+partial def withSemiHilbert {α : Type _}
+  (K X : Expr) (k : Array Expr → MetaM α) : MetaM α := do
+  let .some ⟨_u,K⟩ ← isTypeQ K | throwError "invalid type {← ppExpr K}"
+  loop K X #[] k
+where
+  loop {u} (K : Q(Type $u)) (X : Expr) (acc : Array Expr) (k : Array Expr → MetaM α) : MetaM α := do
+    let cls ← mkAppM ``SemiHilbert #[K, X]
+    match ← synthInstance? cls with
+    | .some _ => k acc
+    | none => 
+      match X with
+      | .forallE _ I Y _ => 
+        if Y.hasLooseBVars then
+          throwError "can't extend context for dependent type `{← ppExpr X}`"
+        withEnumType I u u (fun acc' => loop K Y (acc ++ acc') k)
+      | .app .. => 
+        if X.isAppOfArity' ``Prod 2 then
+          let X₁ := X.getArg! 0
+          let X₂ := X.getArg! 1
+          loop K X₁ acc (fun acc' => loop K X₂ acc' k)
+        else
+          throwError "dont' know how to extend context for the type `{← ppExpr X}`"
+      | .fvar _ => 
+        -- try to upgrade Vec or SemiInnerProductSpace to SemiHilbert
+        let vecX ← mkAppM ``Vec #[K,X]
+        let innerX ← mkAppM ``SemiInnerProductSpace #[K,X]
+        let lctx ← getLCtx
+        let id? ← lctx.findDeclM? 
+          (fun decl => do 
+            pure <| if (← isDefEq decl.type vecX) || (← isDefEq decl.type innerX) 
+                    then .some (decl.fvarId) else .none)
+        match id? with
+        | .some id => 
+            let lctx' := lctx.modifyLocalDecl id
+              fun decl => decl.setType cls
+            let insts ← getLocalInstances
+            let insts' := insts.erase id 
+              |>.push {className := ``SemiHilbert, fvar := .fvar id}
+            withLCtx lctx' insts' (k acc)
+        | .none => 
+          withLocalDecl (← mkAuxName `inst 0) .instImplicit cls fun inst => 
+            k (acc.push inst)
+      | _ => 
+        throwError "dont' know how to extend context for the type `{← ppExpr X}`"
+
+
+
diff --git a/SciLean/Core/Meta/GenerateBasic.lean b/SciLean/Core/Meta/GenerateBasic.lean
index 14325b79..b27feb83 100644
--- a/SciLean/Core/Meta/GenerateBasic.lean
+++ b/SciLean/Core/Meta/GenerateBasic.lean
@@ -149,3 +149,45 @@ def checkNoDependentTypes (xs ys : Array Expr) : MetaM Unit := do
     let Y ← inferType y
     if let .some x := xs.find? (fun x => Y.containsFVar x.fvarId!) then
       throwError s!"the type of ({← ppExpr y} : {← ppExpr (← inferType y)}) depends on the argument {← ppExpr x}, dependent types like this are not supported"
+
+structure GenerateData where
+  /-- field over which we are currently working -/
+  K : Expr
+  
+  /-- original context fvars of a function, these are types, instances and implicit arguments -/
+  orgCtx : Array Expr 
+  /-- extended orgCtx such that types form appropriate vector space, group or whatever is necessary -/
+  ctx : Array Expr 
+
+  /-- main fvars, main arguments we perform function transformation in -/
+  mainArgs : Array Expr
+  /-- unused fvars -/
+  unusedArgs : Array Expr
+  /-- trailing fvars -/
+  trailingArgs : Array Expr
+  /-- argument kinds, this allows to glue arguments back together with mergeArgs and mergeArgs' -/
+  argKinds : Array ArgKind
+
+  /-- names of main arguments guaranteed to be in the same order as mainArgs -/
+  mainNames : Array Name
+
+  /-- auxiliary type we perform transformation in -/
+  W : Expr
+  /-- fvar of type W -/
+  w : Expr
+  /-- fvars making W into vector space, group, or what ever is necessary -/
+  ctxW : Array Expr
+
+  /-- function we are working with as a function of `w` -/
+  f : Expr
+ 
+  /-- fvars that that are main arguments parametrized by W-/
+  argFuns : Array Expr
+  /-- fvars for properties about argFun -/
+  argFunProps : Array Expr
+
+  /-- declaration suffix based on argument names used to generate rule name -/
+  declSuffix : String
+  
+  /-- level parameters -/
+  levelParams : List Name
diff --git a/SciLean/Core/Meta/GenerateRevCDeriv.lean b/SciLean/Core/Meta/GenerateRevCDeriv.lean
index bf535518..47502681 100644
--- a/SciLean/Core/Meta/GenerateRevCDeriv.lean
+++ b/SciLean/Core/Meta/GenerateRevCDeriv.lean
@@ -5,50 +5,9 @@ namespace SciLean.Meta
 
 open Lean Meta Qq
 
-namespace GenerateProperty
+namespace GenerateRevCDeriv
 
-
-structure GenerateData where
-  /-- field over which we are currently working -/
-  K : Expr
-  
-  /-- original context fvars of a function, these are types, instances and implicit arguments -/
-  orgCtx : Array Expr 
-  /-- extended orgCtx such that types form appropriate vector space, group or whatever is necessary -/
-  ctx : Array Expr 
-
-  /-- main fvars, main arguments we perform function transformation in -/
-  mainArgs : Array Expr
-  /-- unused fvars -/
-  unusedArgs : Array Expr
-  /-- trailing fvars -/
-  trailingArgs : Array Expr
-  /-- argument kinds, this allows to glue arguments back together with mergeArgs and mergeArgs' -/
-  argKinds : Array ArgKind
-
-  /-- names of main arguments guaranteed to be in the same order as mainArgs -/
-  mainNames : Array Name
-
-  /-- auxiliary type we perform transformation in -/
-  W : Expr
-  /-- fvar of type W -/
-  w : Expr
-  /-- fvars making W into vector space, group, or what ever is necessary -/
-  ctxW : Array Expr
-
-  /-- function we are working with as a function of `w` -/
-  f : Expr
- 
-  /-- fvars that that are main arguments parametrized by W-/
-  argFuns : Array Expr
-  /-- fvars for properties about argFun -/
-  argFunProps : Array Expr
-
-  /-- declaration suffix based on argument names used to generate rule name -/
-  declSuffix : String
-  
-  /-- level parameters -/
-  levelParams : List Name
+open GenerateProperty
 
 
 /-- Introduce new fvars such that the type `type` have instance of `SemiInnerProductSpace K ·` -/
@@ -213,11 +172,11 @@ def eliminateTransArgFun (e : Expr) (argFuns transArgFuns transArgFunVars : Arra
 
   return e'
 
-end GenerateProperty
+end GenerateRevCDeriv
 
 open Lean Elab Term
 
-open GenerateProperty
+open GenerateRevCDeriv GenerateProperty
 
 def generateRevCDeriv (constName : Name) (mainNames trailingNames : Array Name) (conv : TSyntax `conv) : TermElabM Unit := do
   
diff --git a/test/extend_context.lean b/test/extend_context.lean
new file mode 100644
index 00000000..1d7a83a2
--- /dev/null
+++ b/test/extend_context.lean
@@ -0,0 +1,46 @@
+import SciLean.Core.Meta.ExtendContext
+
+open Lean Meta Qq
+
+open SciLean
+
+
+set_option pp.universes true
+
+/--
+info: newly introduced instances #[SciLean.EnumType.{u, w, w} I, SciLean.SemiInnerProductSpace.{w, v} K X]
+-/
+#guard_msgs in
+#eval show MetaM Unit from 
+  let u := Level.param `u
+  let v := Level.param `v
+  let w := Level.param `w
+  withLocalDeclQ `K .default q(Type $w) fun K => do
+  withLocalDeclQ `r .instImplicit q(IsROrC $K) fun _r => do
+  withLocalDeclQ `I .default q(Type $u) fun I => do
+  withLocalDeclQ `X .default q(Type $v) fun X => do
+    let T := q(($I × $I) → ($I ⊕ $I) → $X)
+
+    withSemiInnerProductSpace K T fun xs => do
+      IO.println s!"newly introduced instances {← xs.mapM (fun x => inferType x >>= ppExpr)} "
+
+
+/--
+info: newly introduced instances #[SciLean.EnumType.{u, w, w} I, SciLean.SemiInnerProductSpace.{w, v} K X, SciLean.EnumType.{u, w, w} J]
+-/
+#guard_msgs in
+#eval show MetaM Unit from 
+  let u := Level.param `u
+  let v := Level.param `v
+  let w := Level.param `w
+  withLocalDeclQ `K .default q(Type $w) fun K => do
+  withLocalDeclQ `r .instImplicit q(IsROrC $K) fun _r => do
+  withLocalDeclQ `I .default q(Type $u) fun I => do
+  withLocalDeclQ `J .default q(Type $u) fun J => do
+  withLocalDeclQ `X .default q(Type $v) fun X => do
+    let T := q((($I × $I) → ($I ⊕ $I) → $X) × ($J → $X))
+
+    withSemiInnerProductSpace K T fun xs => do
+      IO.println s!"newly introduced instances {← xs.mapM (fun x => inferType x >>= ppExpr)} "
+
+