From eb0374dfb3cd73df86bfcfe5a0ed1f5d03d37cce Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Mon, 9 Dec 2024 20:57:25 -0800
Subject: [PATCH] feat: Collect new constants in repl

---
 Pantograph/Protocol.lean |  3 +++
 Repl.lean                | 13 +++++++++++--
 Test/Integration.lean    |  2 ++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/Pantograph/Protocol.lean b/Pantograph/Protocol.lean
index 5cb24e8..d80c761 100644
--- a/Pantograph/Protocol.lean
+++ b/Pantograph/Protocol.lean
@@ -312,6 +312,8 @@ structure FrontendProcess where
   invocations: Bool := false
   -- If set to true, collect `sorry`s
   sorrys: Bool := false
+  -- If set to true, extract new constants
+  newConstants: Bool := false
   deriving Lean.FromJson
 structure InvokedTactic where
   goalBefore: String
@@ -332,6 +334,7 @@ structure CompilationUnit where
   -- Code segments which generated the goals
   goalSrcBoundaries: Array (Nat × Nat) := #[]
   messages: Array String := #[]
+  newConstants: Option (Array String) := .none
   deriving Lean.ToJson
 structure FrontendProcessResult where
   units: List CompilationUnit
diff --git a/Repl.lean b/Repl.lean
index 201a841..710b132 100644
--- a/Repl.lean
+++ b/Repl.lean
@@ -262,9 +262,17 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
           else
             pure []
         let messages ← step.messageStrings
-        return (step.before, boundary, invocations?, sorrys, messages)
+        let newConstants ← if args.newConstants then
+            Frontend.collectNewDefinedConstants step
+          else
+            pure []
+        return (step.before, boundary, invocations?, sorrys, messages, newConstants)
       let li ← frontendM.run context |>.run' state
-      let units ← li.mapM λ (env, boundary, invocations?, sorrys, messages) => Lean.withEnv env do
+      let units ← li.mapM λ (env, boundary, invocations?, sorrys, messages, newConstants) => Lean.withEnv env do
+        let newConstants := if args.newConstants then
+            .some $ newConstants.toArray.map λ name => name.toString
+          else
+            .none
         let (goalStateId?, goals, goalSrcBoundaries) ← if sorrys.isEmpty then do
             pure (.none, #[], #[])
           else do
@@ -279,6 +287,7 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
           goals,
           goalSrcBoundaries,
           messages,
+          newConstants,
         }
       return .ok { units }
     catch e =>
diff --git a/Test/Integration.lean b/Test/Integration.lean
index 171bc91..8d53168 100644
--- a/Test/Integration.lean
+++ b/Test/Integration.lean
@@ -174,6 +174,7 @@ def test_frontend_process : Test :=
         ("file", .str file),
         ("invocations", .bool true),
         ("sorrys", .bool false),
+        ("newConstants", .bool false),
       ]
      ({
        units := [{
@@ -214,6 +215,7 @@ def test_frontend_process_sorry : Test :=
         ("file", .str file),
         ("invocations", .bool false),
         ("sorrys", .bool true),
+        ("newConstants", .bool false),
       ]
      ({
        units := [{