From b9ff9e8f136fb2354b3e5262658d114602a2a949 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Sat, 29 Mar 2025 14:51:35 -0700
Subject: [PATCH] feat(repl): Optional type argument in `env.add`

---
 Pantograph/Environment.lean | 19 ++++++++++++-------
 Pantograph/Protocol.lean    |  2 +-
 Repl.lean                   |  2 +-
 Test/Integration.lean       |  5 ++---
 4 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/Pantograph/Environment.lean b/Pantograph/Environment.lean
index d7bac21..21d7b37 100644
--- a/Pantograph/Environment.lean
+++ b/Pantograph/Environment.lean
@@ -144,29 +144,34 @@ def inspect (args: Protocol.EnvInspect) (options: @&Protocol.Options): Protocol.
     else
       .pure result
   return result
+/-- Elaborates and adds a declaration to the `CoreM` environment. -/
 @[export pantograph_env_add_m]
-def addDecl (name: String) (levels: Array String := #[]) (type: String) (value: String) (isTheorem: Bool)
+def addDecl (name: String) (levels: Array String := #[]) (type?: Option String) (value: String) (isTheorem: Bool)
     : Protocol.FallibleT CoreM Protocol.EnvAddResult := do
   let env ← Lean.MonadEnv.getEnv
   let levelParams := levels.toList.map (·.toName)
   let tvM: Elab.TermElabM (Except String (Expr × Expr)) :=
     Elab.Term.withLevelNames levelParams do do
-    let type ← match parseTerm env type with
-      | .ok syn => do
-        match ← elabTerm syn with
-        | .error e => return .error e
-        | .ok expr => pure expr
+    let expectedType?? : Except String (Option Expr) ← ExceptT.run $ type?.mapM λ type => do
+      match parseTerm env type with
+      | .ok syn => elabTerm syn
+      | .error e => MonadExceptOf.throw e
+    let expectedType? ← match expectedType?? with
+      | .ok t? => pure t?
       | .error e => return .error e
     let value ← match parseTerm env value with
       | .ok syn => do
         try
-          let expr ← Elab.Term.elabTerm (stx := syn) (expectedType? := .some type)
+          let expr ← Elab.Term.elabTerm (stx := syn) (expectedType? := expectedType?)
           Lean.Elab.Term.synthesizeSyntheticMVarsNoPostponing
           let expr ← instantiateMVars expr
           pure $ expr
         catch ex => return .error (← ex.toMessageData.toString)
       | .error e => return .error e
     Elab.Term.synthesizeSyntheticMVarsNoPostponing
+    let type ← match expectedType? with
+      | .some t => pure t
+      | .none => Meta.inferType value
     pure $ .ok (← instantiateMVars type, ← instantiateMVars value)
   let (type, value) ← match ← tvM.run' (ctx := {}) |>.run' with
     | .ok t => pure t
diff --git a/Pantograph/Protocol.lean b/Pantograph/Protocol.lean
index 41a6ccd..350b1de 100644
--- a/Pantograph/Protocol.lean
+++ b/Pantograph/Protocol.lean
@@ -205,7 +205,7 @@ structure EnvInspectResult where
 structure EnvAdd where
   name: String
   levels: Array String := #[]
-  type: String
+  type?: Option String := .none
   value: String
   isTheorem: Bool := false
   deriving Lean.FromJson
diff --git a/Repl.lean b/Repl.lean
index 3c87429..578bd51 100644
--- a/Repl.lean
+++ b/Repl.lean
@@ -227,7 +227,7 @@ def execute (command: Protocol.Command): MainM Json := do
     let state ← getMainState
     runCoreM' $ Environment.inspect args state.options
   env_add (args: Protocol.EnvAdd): EMainM Protocol.EnvAddResult := do
-    runCoreM' $ Environment.addDecl args.name args.levels args.type args.value args.isTheorem
+    runCoreM' $ Environment.addDecl args.name args.levels args.type? args.value args.isTheorem
   env_save (args: Protocol.EnvSaveLoad): EMainM Protocol.EnvSaveLoadResult := do
     let env ← MonadEnv.getEnv
     environmentPickle env args.path
diff --git a/Test/Integration.lean b/Test/Integration.lean
index 4778f0a..9c3bc5e 100644
--- a/Test/Integration.lean
+++ b/Test/Integration.lean
@@ -167,7 +167,6 @@ def test_env_add_inspect : Test :=
     step "env.add"
       ({
         name := name1,
-        type := "Prop → Prop → Prop",
         value := "λ (a b: Prop) => Or a b",
         isTheorem := false
       }: Protocol.EnvAdd)
@@ -181,7 +180,7 @@ def test_env_add_inspect : Test :=
     step "env.add"
       ({
         name := name2,
-        type := "Nat → Int",
+        type? := "Nat → Int",
         value := "λ (a: Nat) => a + 1",
         isTheorem := false
       }: Protocol.EnvAdd)
@@ -196,7 +195,7 @@ def test_env_add_inspect : Test :=
       ({
         name := name3,
         levels := #["u"]
-        type := "(α : Type u) → α → (α × α)",
+        type? := "(α : Type u) → α → (α × α)",
         value := "λ (α : Type u) (x : α) => (x, x)",
         isTheorem := false
       }: Protocol.EnvAdd)