Add opt-in declaration and whole-proof profiling

Beam/Broker/Protocol.lean

@@ -264,6 +264,7 @@ structure CloseRequest extends RequestFile where
 structure RunAtRequest extends RequestPosition where
   text : String
   storeHandle? : Option Bool := none
+  profile? : Option Bool := none
 
 structure ReferencesRequest extends RequestPosition where
   includeDeclaration? : Option Bool := none
@@ -294,6 +295,7 @@ structure RunWithRequest where
   text : String
   storeHandle? : Option Bool := none
   linear? : Option Bool := none
+  profile? : Option Bool := none
   handle : Handle
 
 structure ReleaseRequest where
@@ -433,7 +435,7 @@ private def Op.requestFields (op : Op) : Array String :=
       #["path", "diagnosticScope", "diagnosticsInResult"]
   | .close => #["path", "diagnosticScope", "saveArtifacts"]
   | .runAt =>
-      #["path", "version", "line", "character", "text", "storeHandle"]
+      #["path", "version", "line", "character", "text", "storeHandle", "profile"]
   | .hover | .signatureHelp | .definition =>
       #["path", "version", "line", "character"]
   | .references =>
@@ -453,7 +455,7 @@ private def Op.requestFields (op : Op) : Array String :=
         "suggest"
       ]
   | .runWith =>
-      #["path", "text", "storeHandle", "linear", "handle"]
+      #["path", "text", "storeHandle", "linear", "profile", "handle"]
   | .release => #["path", "handle"]
   | .initWorkspace =>
       #["workspaceMode", "root", "leanCmd", "leanPlugin", "rocqCmd"]
@@ -501,7 +503,8 @@ private def RequestPayload.jsonFields : RequestPayload → List (String × Json)
   | .runAt request =>
       request.toRequestPosition.jsonFields ++
       [("text", toJson request.text)] ++
-      optionalJsonField "storeHandle" request.storeHandle?
+      optionalJsonField "storeHandle" request.storeHandle? ++
+      optionalJsonField "profile" request.profile?
   | .hover request | .signatureHelp request | .definition request =>
       request.jsonFields
   | .references request =>
@@ -532,6 +535,7 @@ private def RequestPayload.jsonFields : RequestPayload → List (String × Json)
       [("path", toJson request.path), ("text", toJson request.text)] ++
       optionalJsonField "storeHandle" request.storeHandle? ++
       optionalJsonField "linear" request.linear? ++
+      optionalJsonField "profile" request.profile? ++
       [("handle", toJson request.handle)]
   | .release request =>
       [("path", toJson request.path), ("handle", toJson request.handle)]
@@ -681,6 +685,7 @@ instance : FromJson Request where
             toRequestPosition := target
             text := ← requiredField j "text"
             storeHandle? := ← optionalField? (α := Bool) j "storeHandle"
+            profile? := ← optionalField? (α := Bool) j "profile"
           }
       | .hover | .signatureHelp | .definition => do
           let target ← decodeRequestPosition j backend
@@ -737,6 +742,7 @@ instance : FromJson Request where
             text := ← requiredField j "text"
             storeHandle? := ← optionalField? (α := Bool) j "storeHandle"
             linear? := ← optionalField? (α := Bool) j "linear"
+            profile? := ← optionalField? (α := Bool) j "profile"
             handle
           }
       | .release => do

Beam/Broker/Server.lean

@@ -2212,10 +2212,11 @@ private def handleRunAtOp
         [ ("textDocument", toJson ({ uri := uri, version? := some request.version : VersionedTextDocumentIdentifier }))
         , ("position", toJson ({ line := request.line, character := request.character : Lsp.Position }))
         , ("text", toJson request.text)
-        ] ++
-        match request.storeHandle? with
+        ] ++ (match request.storeHandle? with
         | some b => [("storeHandle", toJson b)]
-        | none => [])
+        | none => []) ++ (match request.profile? with
+        | some b => [("profile", toJson b)]
+        | none => []))
       trackedDocumentVersion
       (expectedVersion? := some request.version)
       (clientRequestId? := req.clientRequestId?)
@@ -2516,6 +2517,9 @@ private def handleRunWithOp
             | none => []) ++
             (match request.linear? with
             | some b => [("linear", toJson b)]
+            | none => []) ++
+            (match request.profile? with
+            | some b => [("profile", toJson b)]
             | none => []))
           trackedDocumentVersion
           (clientRequestId? := req.clientRequestId?)

Beam/Cli/Args.lean

@@ -65,6 +65,14 @@ def parseTextArg (cmdHead : String) (args : List String) : IO ParsedTextArg := d
         | throw <| IO.userError (textArgUsage cmdHead)
       pure { text, source := "argv" }
 
+/-- Consume the optional profiling selector before a text input selector. -/
+def parseProfileArg (args : List String) : IO (Bool × List String) := do
+  match args with
+  | "--profile" :: "--profile" :: _ =>
+      throw <| IO.userError "duplicate --profile"
+  | "--profile" :: rest => pure (true, rest)
+  | _ => pure (false, args)
+
 def parseJsonText (label text : String) : IO Json := do
   match Json.parse text with
   | .ok json => pure json

Beam/Cli/Commands.lean

@@ -96,11 +96,13 @@ private def runLeanRunAt
   let version ← parseNatArg "version" versionText
   let line ← parseNatArg "line" lineText
   let character ← parseNatArg "character" characterText
+  let (profile, textArgs) ← parseProfileArg textArgs
   let parsedText ← parseTextArg s!"{action} <path> <version> <line> <character>" textArgs
   let root ← projectRoot opts .lean
   withProjectDaemon root .lean (explicitControlDir? := opts.explicitControlDir?) fun client => do
     let req ← withEnvClientRequestId <|
-      leanRunAtRequest path version line character parsedText.text (storeHandle := storeHandle)
+      leanRunAtRequest path version line character parsedText.text
+        (storeHandle := storeHandle) (profile := profile)
     maybeEmitTextDebug req.clientRequestId? action parsedText.source parsedText.text
     callBrokerWithProgress root client req (leanRunAtWaitSpec action path line character)
 
@@ -114,16 +116,18 @@ private def runLeanRunWith
     | [] => []
     | "--handle-file" :: _ :: rest => rest
     | _ :: rest => rest
-  if handleArgReadsStdin args && textArgReadsStdin textArgs then
+  let (_, profileFreeTextArgs) ← parseProfileArg textArgs
+  if handleArgReadsStdin args && textArgReadsStdin profileFreeTextArgs then
     throw <| IO.userError <| String.intercalate "\n" [
       textArgUsage s!"{action} <path> <handle-json|-|--handle-file <path>>",
       "cannot read both handle json and continuation text from stdin; pass the handle inline, use --handle-file, or use --text-file for the text"
     ]
   let (handle, textArgs) ← parseHandleInput s!"{action} <path>" args
+  let (profile, textArgs) ← parseProfileArg textArgs
   let parsedText ← parseTextArg s!"{action} <path> <handle-json|-|--handle-file <path>>" textArgs
   let root ← projectRoot opts .lean
   let req ← withEnvClientRequestId <|
-    leanRunWithRequest path handle parsedText.text (linear := linear)
+    leanRunWithRequest path handle parsedText.text (linear := linear) (profile := profile)
   maybeEmitTextDebug req.clientRequestId? action parsedText.source parsedText.text
   withProjectDaemon root .lean (explicitControlDir? := opts.explicitControlDir?) fun client =>
     callBrokerWithProgress root client req (leanRunWithWaitSpec path (linear := linear))

Beam/Cli/LeanOperation.lean

@@ -17,16 +17,18 @@ def leanRunAtRequest
     (version : Nat)
     (line character : Nat)
     (text : String)
-    (storeHandle : Bool := false) : Request :=
-  ({ path, version, line, character, text } : Beam.Lean.RunAtInput).toBrokerRequest
+    (storeHandle : Bool := false)
+    (profile : Bool := false) : Request :=
+  ({ path, version, line, character, text, profile? := if profile then some true else none } : Beam.Lean.RunAtInput).toBrokerRequest
     (storeHandle := storeHandle)
 
 def leanRunWithRequest
     (path : String)
     (handle : Handle)
     (text : String)
-    (linear : Bool := false) : Request :=
-  ({ path, handle, text } : Beam.Lean.RunWithInput).toBrokerRequest
+    (linear : Bool := false)
+    (profile : Bool := false) : Request :=
+  ({ path, handle, text, profile? := if profile then some true else none } : Beam.Lean.RunWithInput).toBrokerRequest
     (linear := linear)
 
 def leanReleaseRequest (path : String) (handle : Handle) : Request :=

Beam/Cli/Usage.lean

@@ -14,8 +14,8 @@ def usage : String :=
     "  lean-beam --version",
     "  lean-beam version",
     "  lean-beam [--root PATH] [--session-dir DIR] serve [lean|rocq]",
-    "  lean-beam [--root PATH] run-at <path> <version> <line> <character> (--stdin | --text-file <path> | -- <text...> | <text...>)",
-    "  lean-beam [--root PATH] run-at-handle <path> <version> <line> <character> (--stdin | --text-file <path> | -- <text...> | <text...>)",
+    "  lean-beam [--root PATH] run-at <path> <version> <line> <character> [--profile] (--stdin | --text-file <path> | -- <text...> | <text...>)",
+    "  lean-beam [--root PATH] run-at-handle <path> <version> <line> <character> [--profile] (--stdin | --text-file <path> | -- <text...> | <text...>)",
     "  lean-beam [--root PATH] hover <path> <version> <line> <character>",
     "  lean-beam [--root PATH] signature-help <path> <version> <line> <character>",
     "  lean-beam [--root PATH] definition <path> <version> <line> <character>",
@@ -24,8 +24,8 @@ def usage : String :=
     "  lean-beam [--root PATH] workspace-symbols <query...>",
     "  lean-beam [--root PATH] goals before|after <path> <version> <line> <character>",
     "  lean-beam [--root PATH] todo <path> <version> <startLine> <startCharacter> <endLine> <endCharacter> [--kind <kind> ...] [--suggest none|basic]",
-    "  lean-beam [--root PATH] run-with <path> <handle-json|-|--handle-file <path>> (--stdin | --text-file <path> | -- <text...> | <text...>)",
-    "  lean-beam [--root PATH] run-with-linear <path> <handle-json|-|--handle-file <path>> (--stdin | --text-file <path> | -- <text...> | <text...>)",
+    "  lean-beam [--root PATH] run-with <path> <handle-json|-|--handle-file <path>> [--profile] (--stdin | --text-file <path> | -- <text...> | <text...>)",
+    "  lean-beam [--root PATH] run-with-linear <path> <handle-json|-|--handle-file <path>> [--profile] (--stdin | --text-file <path> | -- <text...> | <text...>)",
     "  lean-beam [--root PATH] release <path> <handle-json|-|--handle-file <path>>",
     "  lean-beam [--root PATH] update <path>",
     "  lean-beam [--root PATH] sync <path> [+all-diagnostics]",