Let --keep-all-files emit all intermediate representations after each pass of Laurel and Core (#840)

This patch makes `--keep-all-files` emit all intermediate
representations after each pass of Laurel and Core. :)

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Author: aqjune-aws

Strata/Languages/Core/Verifier.lean

@@ -935,29 +935,47 @@ def verifySingleEnv (pE : Program × Env) (options : VerifyOptions)
 /-- Run the Strata Core verification pipeline on a program: transform,
 type-check, partially evaluate, and discharge proof obligations via SMT.
 All program-wide transformations that occur before any analyses
-(including type inference) should be placed here. -/
+(including type inference) should be placed here.
+
+When `keepAllFilesPrefix` is provided, the program state after each pipeline
+phase is written to `{prefix}.{n}.{phaseName}.core.st` (numbered from 1). -/
 def verify (program : Program)
     (tempDir : System.FilePath)
     (proceduresToVerify : Option (List String) := none)
     (options : VerifyOptions := VerifyOptions.default)
     (moreFns : @Lambda.Factory CoreLParams := Lambda.Factory.default)
     (externalPhases : List AbstractedPhase := [])
     (prefixPhases : List PipelinePhase := [])
+    (keepAllFilesPrefix : Option String := none)
     : EIO DiagnosticModel VCResults := do
   let profile := options.profile
   let factory ← EIO.ofExcept (Core.Factory.addFactory moreFns)
   let pipelinePhases := prefixPhases ++ corePipelinePhases (procs := proceduresToVerify) (factory := some factory)
   let phases := pipelinePhases.map (·.phase)
   let finalProgram ← profileStep profile "  Program transformations" do
-    let passes := fun prog => do
-      let mut current := prog
-      for pp in pipelinePhases do
-        let (_changed, next) ← pp.transform current
+    if let some pfx := keepAllFilesPrefix then
+      if let some parent := (System.FilePath.mk pfx).parent then
+        IO.toEIO (fun e => DiagnosticModel.fromFormat f!"{e}")
+          (IO.FS.createDirAll parent)
+    let mut current := program
+    let mut state : Transform.CoreTransformState := .emp
+    let mut step := 0
+    for pp in pipelinePhases do
+      let (result, newState) := Transform.runWith current (fun prog => do
+        let (_, next) ← pp.transform prog
+        return next) state
+      match result with
+      | .ok next =>
         current := next
-      return current
-    match Transform.run program passes with
-    | .ok prog => .ok prog
-    | .error e => .error (DiagnosticModel.fromFormat f!"❌ Transform Error. {e}")
+        state := newState
+        step := step + 1
+        if let some pfx := keepAllFilesPrefix then
+          let path := s!"{pfx}.{step}.{pp.phase.name}.core.st"
+          IO.toEIO (fun e => DiagnosticModel.fromFormat f!"{e}")
+            (IO.FS.writeFile path (toString current ++ "\n"))
+      | .error e =>
+        throw (DiagnosticModel.fromFormat f!"❌ Transform Error. {e}")
+    .ok current
   -- Build the axiom relevance cache once (post-transform, so declarations are
   -- stable). The cache is reused across all verification environments and goals.
   let axiomCache? ← profileStep profile "  Build axiom relevance cache" do
@@ -992,7 +1010,6 @@ def typeCheck (ictx : InputContext) (env : Program) (options : Core.VerifyOption
   Except DiagnosticModel Core.Program := do
   let (program, errors) := TransM.run ictx (translateProgram env)
   if errors.isEmpty then
-    -- dbg_trace f!"AST: {program}"
     Core.typeCheck options program moreFns
   else
     .error <| DiagnosticModel.fromFormat s!"DDM Transform Error: {repr errors}"
@@ -1017,13 +1034,15 @@ def verify
     (options : Core.VerifyOptions := Core.VerifyOptions.default)
     (moreFns : @Lambda.Factory Core.CoreLParams := Lambda.Factory.default)
     (externalPhases : List Core.AbstractedPhase := [])
+    (keepAllFilesPrefix : Option String := none)
     : IO Core.VCResults := do
   let (program, errors) := Core.getProgram env ictx
   if errors.isEmpty then
     let runner tempDir :=
       EIO.toIO (fun dm => IO.Error.userError (toString (dm.format (some ictx.fileMap))))
                   (Core.verify program tempDir proceduresToVerify options moreFns
-                    (externalPhases := externalPhases))
+                    (externalPhases := externalPhases)
+                    (keepAllFilesPrefix := keepAllFilesPrefix))
     match options.vcDirectory with
     | .none =>
       IO.FS.withTempDir runner

Strata/Languages/Laurel/LaurelToCoreTranslator.lean

@@ -699,51 +699,79 @@ abbrev TranslateResultWithLaurel := (Option Core.Program) × (List DiagnosticMod
 
 /--
 Translate Laurel Program to Core Program, also returning the lowered Laurel program.
+
+When `keepAllFilesPrefix` is provided, the program state after each named
+Laurel-to-Laurel pass is written to `{prefix}.{n}.{passName}.laurel.st`
+(numbered from 1).
 -/
-def translateWithLaurel (options: LaurelTranslateOptions) (program : Program): TranslateResultWithLaurel :=
+def translateWithLaurel (options: LaurelTranslateOptions) (program : Program)
+    (keepAllFilesPrefix : Option String := none)
+    : IO TranslateResultWithLaurel := do
   let program := { program with
     staticProcedures := coreDefinitionsForLaurel.staticProcedures ++ program.staticProcedures
   }
+  if let some pfx := keepAllFilesPrefix then
+    if let some parent := (System.FilePath.mk pfx).parent then
+      IO.FS.createDirAll parent
+  let stepRef ← IO.mkRef (0 : Nat)
+  let emit (name : String) (p : Program) : IO Unit :=
+    match keepAllFilesPrefix with
+    | some pfx => do
+      let n ← stepRef.modifyGet (fun n => (n, n + 1))
+      IO.FS.writeFile s!"{pfx}.{n}.{name}.laurel.st"
+        ((formatProgram p).pretty ++ "\n")
+    | none => pure ()
+
+  -- Step 0: the input program before any passes
+  emit "Initial" program
 
-  -- dbg_trace "=== Initial Laurel program ==="
-  -- dbg_trace (toString (Std.Format.pretty (Std.ToFormat.format program)))
-  -- dbg_trace "================================="
   let result := resolve program
   let resolutionErrors: List DiagnosticModel := if options.emitResolutionErrors then result.errors.toList else []
   let (program, model) := (result.program, result.model)
+  emit "Resolve" program
   let diamondErrors := validateDiamondFieldAccesses model program
 
   let (program, nonCompositeDiags) := filterNonCompositeModifies model program
+  emit "FilterNonCompositeModifies" program
 
   let program := heapParameterization model program
   let result := resolve program (some model)
   let (program, model) := (result.program, result.model)
+  emit "HeapParameterization" program
 
   let program := typeHierarchyTransform model program
   let result := resolve program (some model)
   let (program, model) := (result.program, result.model)
+  emit "TypeHierarchyTransform" program
   let (program, modifiesDiags) := modifiesClausesTransform model program
   let result := resolve program (some model)
   let (program, model) := (result.program, result.model)
   let result := resolve program (some model)
   let (program, model) := (result.program, result.model)
+  emit "ModifiesClausesTransform" program
   let program := inferHoleTypes model program
+  emit "InferHoleTypes" program
   let program := eliminateHoles program
+  emit "EliminateHoles" program
   let program := desugarShortCircuit model program
+  emit "DesugarShortCircuit" program
   let program := liftExpressionAssignments model program
+  emit "LiftExpressionAssignments" program
   let program := eliminateReturnsInExpressionTransform program
   let result := resolve program (some model)
   let (program, model) := (result.program, result.model)
+  emit "EliminateReturns" program
 
   let (program, constrainedTypeDiags) := constrainedTypeElim model program
   let result := resolve program (some model)
   let (program, model) := (result.program, result.model)
+  emit "ConstrainedTypeElim" program
 
   let initState : TranslateState := {model := model }
   let (coreProgramOption, translateState) := runTranslateM initState (translateLaurelToCore options program)
   let allDiagnostics := resolutionErrors ++ diamondErrors ++ nonCompositeDiags ++ modifiesDiags ++ constrainedTypeDiags ++ translateState.diagnostics
   let coreProgramOption := if translateState.coreProgramHasSuperfluousErrors then none else coreProgramOption
-  (coreProgramOption, allDiagnostics, program)
+  return (coreProgramOption, allDiagnostics, program)
   where
 
   /--
@@ -824,17 +852,17 @@ def translateWithLaurel (options: LaurelTranslateOptions) (program : Program): T
 /--
 Translate Laurel Program to Core Program
 -/
-def translate (options: LaurelTranslateOptions) (program : Program): TranslateResult :=
-  let (core, diags, _) := translateWithLaurel options program
-  (core, diags)
+def translate (options: LaurelTranslateOptions) (program : Program): IO TranslateResult := do
+  let (core, diags, _) ← translateWithLaurel options program
+  return (core, diags)
 
 /--
 Verify a Laurel program using an SMT solver
 -/
 def verifyToVcResults (program : Program)
     (options : VerifyOptions := .default)
     : IO (Option VCResults × List DiagnosticModel) := do
-  let (coreProgramOption, translateDiags) := translate {} program
+  let (coreProgramOption, translateDiags) ← translate {} program
 
   match coreProgramOption with
   | some coreProgram =>

Strata/Languages/Python/PySpecPipeline.lean

@@ -359,18 +359,24 @@ public def splitProcNames (prog : Core.Program)
   (preludeNames, userProcNames)
 
 /-- Like `translateCombinedLaurel` but also returns the lowered Laurel program
-    (after all Laurel-to-Laurel passes, before translation to Core). -/
+    (after all Laurel-to-Laurel passes, before translation to Core).
+
+    When `keepAllFilesPrefix` is provided, the program state after each named
+    Laurel pass is written to `{prefix}.{n}.{passName}.laurel.st`. -/
 public def translateCombinedLaurelWithLowered (combined : Laurel.Program)
-    : (Option Core.Program × List DiagnosticModel × Laurel.Program) :=
-  let (coreOption, errors, lowered) := Laurel.translateWithLaurel { inlineFunctionsWhenPossible := true } combined
-  (coreOption.map appendCorePartOfRuntime, errors, lowered)
+    (keepAllFilesPrefix : Option String := none)
+    : IO (Option Core.Program × List DiagnosticModel × Laurel.Program) := do
+  let (coreOption, errors, lowered) ←
+    Laurel.translateWithLaurel { inlineFunctionsWhenPossible := true } combined
+      (keepAllFilesPrefix := keepAllFilesPrefix)
+  return (coreOption.map appendCorePartOfRuntime, errors, lowered)
 
 /-- Translate a combined Laurel program to Core and prepend the full
     runtime prelude. -/
 public def translateCombinedLaurel (combined : Laurel.Program)
-    : (Option Core.Program × List DiagnosticModel) :=
-  let (coreOption, errors, _) := translateCombinedLaurelWithLowered combined
-  (coreOption, errors)
+    : IO (Option Core.Program × List DiagnosticModel) := do
+  let (coreOption, errors, _) ← translateCombinedLaurelWithLowered combined
+  return (coreOption, errors)
 
 /-- Errors from the pyAnalyzeLaurel pipeline. -/
 public inductive PipelineError where
@@ -399,7 +405,7 @@ public instance : ToString PipelineError where
     The optional `sourcePath` overrides the file path embedded in
     Laurel metadata (useful when the Ion file was generated from a
     `.py` source and you want line numbers to refer to the original). -/
-public def pyAnalyzeLaurel
+public def pythonAndSpecToLaurel
     (pythonIonPath : String)
     (dispatchModules : Array String := #[])
     (pyspecModules : Array String := #[])

Strata/SimpleAPI.lean

@@ -233,11 +233,11 @@ dialect into the dialect-specific AST for the Core dialect. This can fail with
 an error message if the input program contains constructs that are not yet
 supported.
 -/
-def laurelToCore (p : Laurel.Program) : Except String Core.Program :=
-  let (coreOpt, diags) := Laurel.translate { emitResolutionErrors := true } p
+def laurelToCore (p : Laurel.Program) : IO (Except String Core.Program) := do
+  let (coreOpt, diags) ← Laurel.translate { emitResolutionErrors := true } p
   match coreOpt with
-  | some core => .ok core
-  | none => .error s!"Laurel to Core translation failed: {diags.map (·.message)}"
+  | some core => return .ok core
+  | none => return .error s!"Laurel to Core translation failed: {diags.map (·.message)}"
 
 /-! ### Transformation of Core programs -/
 
@@ -327,10 +327,12 @@ def Core.verifyProgram
     (proceduresToVerify : Option (List String) := none)
     (externalPhases : List Core.AbstractedPhase := [])
     (prefixPhases : List Core.PipelinePhase := [])
+    (keepAllFilesPrefix : Option String := none)
     : EIO String Core.VCResults := do
   let runVerification (tempDir : System.FilePath) : IO Core.VCResults :=
     EIO.toIO (IO.Error.userError ∘ toString)
-      (Core.verify program tempDir proceduresToVerify options moreFns externalPhases prefixPhases)
+      (Core.verify program tempDir proceduresToVerify options moreFns externalPhases prefixPhases
+        (keepAllFilesPrefix := keepAllFilesPrefix))
   let ioAction := match options.vcDirectory with
     | .some vcDir => IO.FS.createDirAll vcDir *> runVerification vcDir
     | .none => IO.FS.withTempDir runVerification
@@ -539,7 +541,7 @@ def pySpecsDir (sourceDir strataDir dialectFile : System.FilePath)
 /-! ### Python-to-Core via Laurel pipeline -/
 
 /-- Translate a Python Ion file all the way to Core.  Composes
-    `pyAnalyzeLaurel` (Python → combined Laurel) and
+    `pythonAndSpecToLaurel` (Python → combined Laurel) and
     `translateCombinedLaurel` (Laurel → Core with prelude). -/
 def pyTranslateLaurel
     (pythonIonPath : String)
@@ -548,10 +550,10 @@ def pyTranslateLaurel
     (specDir : System.FilePath := ".")
     : EIO String (Core.Program × List DiagnosticModel) := do
   let laurel ←
-    match ← pyAnalyzeLaurel pythonIonPath dispatchModules pyspecModules (specDir := specDir) |>.toBaseIO with
+    match ← pythonAndSpecToLaurel pythonIonPath dispatchModules pyspecModules (specDir := specDir) |>.toBaseIO with
     | .ok r => pure r
     | .error err => throw (toString err)
-  let (coreOption, laurelTranslateErrors) := translateCombinedLaurel laurel
+  let (coreOption, laurelTranslateErrors) ← IO.toEIO (fun e => s!"{e}") (translateCombinedLaurel laurel)
   match coreOption with
   | none => throw s!"Laurel to Core translation failed: {laurelTranslateErrors}"
   | some core => pure (core, laurelTranslateErrors)

StrataMain.lean

@@ -524,7 +524,7 @@ def pyAnalyzeLaurelCommand : Command where
       | some (pyPath, srcText) => some (pyPath, .ofString srcText)
       | none => none
     let combinedLaurel ←
-      match ← Strata.pyAnalyzeLaurel filePath dispatchModules pyspecModules sourcePath
+      match ← Strata.pythonAndSpecToLaurel filePath dispatchModules pyspecModules sourcePath
                 (specDir := specDir) (profile := profile)
                 (quiet := quiet) |>.toBaseIO with
       | .ok r => pure r
@@ -558,17 +558,12 @@ def pyAnalyzeLaurelCommand : Command where
       IO.println "\n==== Laurel Program ===="
       IO.println f!"{combinedLaurel}"
 
-    if let some dir := keepDir then
-      let path := s!"{dir}/{baseName}.laurel"
-      IO.FS.writeFile path ((Laurel.formatProgram combinedLaurel).pretty ++ "\n")
+    let keepPrefix := keepDir.map (s!"{·}/{baseName}")
 
-    let (coreProgramOption, laurelTranslateErrors, loweredLaurel) ←
+    let (coreProgramOption, laurelTranslateErrors, _loweredLaurel) ←
       profileStep profile "Laurel to Core translation" do
-        pure (Strata.translateCombinedLaurelWithLowered combinedLaurel)
-
-    if let some dir := keepDir then
-      let path := s!"{dir}/{baseName}.lowered.laurel"
-      IO.FS.writeFile path ((Laurel.formatProgram loweredLaurel).pretty ++ "\n")
+        Strata.translateCombinedLaurelWithLowered combinedLaurel
+          (keepAllFilesPrefix := keepPrefix)
 
     let coreProgram ←
       match coreProgramOption with
@@ -586,10 +581,6 @@ def pyAnalyzeLaurelCommand : Command where
     let (_preludeNames, userProcNames) :=
       Strata.splitProcNames coreProgram [userSourcePath]
 
-    if let some dir := keepDir then
-      let path := s!"{dir}/{baseName}.core"
-      IO.FS.writeFile path (toString coreProgram)
-
     -- Verify using Core verifier
     -- --keep-all-files implies vc-directory if not explicitly set
     let baseVcDir := keepDir.map (fun dir => (s!"{dir}/{baseName}" : System.FilePath))
@@ -611,7 +602,9 @@ def pyAnalyzeLaurelCommand : Command where
                 (moreFns := Strata.Python.ReFactory)
                 (proceduresToVerify := some userProcNames)
                 (externalPhases := [Strata.frontEndPhase])
-                (prefixPhases := inlinePhases) |>.toBaseIO with
+                (prefixPhases := inlinePhases)
+                (keepAllFilesPrefix := keepPrefix)
+                |>.toBaseIO with
       | .ok r => pure r
       | .error msg => exitPyAnalyzeInternalError msg
 
@@ -883,7 +876,7 @@ def laurelAnalyzeToGotoCommand : Command where
     let path : System.FilePath := v[0]
     let content ← IO.FS.readFile path
     let laurelProgram ← Strata.parseLaurelText path content
-    match Strata.Laurel.translate {} laurelProgram with
+    match ← Strata.Laurel.translate {} laurelProgram with
       | (none, diags) => exitFailure s!"Core translation errors: {diags.map (·.message)}"
       | (some coreProgram, errors) =>
         let Ctx := { Lambda.LContext.default with functions := Core.Factory, knownTypes := Core.KnownTypes }
@@ -1008,7 +1001,7 @@ def laurelToCoreCommand : Command where
   help := "Translate a Laurel source file to Core and print to stdout."
   callback := fun v _ => do
     let laurelProgram ← Strata.readLaurelTextFile v[0]
-    let (coreProgramOption, errors) := Strata.Laurel.translate {} laurelProgram
+    let (coreProgramOption, errors) ← Strata.Laurel.translate {} laurelProgram
       if !errors.isEmpty then
         IO.println s!"Core translation errors: {errors.map (·.message)}"
       match coreProgramOption with

StrataTest/Languages/Laurel/Examples/Fundamentals/T8c_BodilessInlining.lean

@@ -32,7 +32,7 @@ procedure caller() {
 #guard_msgs in
 #eval show IO String from do
   let laurelProg ← Strata.parseLaurelText "test.laurel" laurelSource
-  let coreProg ← match Strata.laurelToCore laurelProg with
+  let coreProg ← match ← Strata.laurelToCore laurelProg with
     | .ok p => pure p
     | .error e => throw (IO.userError s!"Translation failed: {e}")
   let inlined ← match Strata.Core.inlineProcedures coreProg {} with