From 82d313bda94af1092aeda02c3618e81fbaf7c9eb Mon Sep 17 00:00:00 2001
From: Andrew Wells <anmwells@amazon.com>
Date: Mon, 1 Dec 2025 10:48:20 -0600
Subject: [PATCH 1/8] S3Client -> Client

---
 Strata/Languages/Python/PythonToBoogie.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Strata/Languages/Python/PythonToBoogie.lean b/Strata/Languages/Python/PythonToBoogie.lean
index 1aed3db7d..d25cd4bd1 100644
--- a/Strata/Languages/Python/PythonToBoogie.lean
+++ b/Strata/Languages/Python/PythonToBoogie.lean
@@ -213,7 +213,7 @@ def collectVarDecls (stmts: Array (Python.stmt SourceRange)) : List Boogie.State
     | "str" => [(.init name t[string] (.strConst () "")), (.havoc name)]
     | "int" => [(.init name t[int] (.intConst () 0)), (.havoc name)]
     | "bytes" => [(.init name t[string] (.strConst () "")), (.havoc name)]
-    | "S3Client" => [(.init name clientType dummyClient), (.havoc name)]
+    | "Client" => [(.init name clientType dummyClient), (.havoc name)]
     | "Dict[str Any]" => [(.init name dictStrAnyType dummyDictStrAny), (.havoc name)]
     | _ => panic! s!"Unsupported type annotation: `{ty_name}`"
   let foo := dedup.map toBoogie

From 31bea64736f900191863fb7c17f683824b2e02fb Mon Sep 17 00:00:00 2001
From: Andrew Wells <anmwells@amazon.com>
Date: Mon, 1 Dec 2025 13:56:03 -0600
Subject: [PATCH 2/8] checkpoint

---
 Strata/Languages/Python/PythonToBoogie.lean | 121 ++++++++++++++------
 1 file changed, 89 insertions(+), 32 deletions(-)

diff --git a/Strata/Languages/Python/PythonToBoogie.lean b/Strata/Languages/Python/PythonToBoogie.lean
index d25cd4bd1..14043a654 100644
--- a/Strata/Languages/Python/PythonToBoogie.lean
+++ b/Strata/Languages/Python/PythonToBoogie.lean
@@ -74,29 +74,51 @@ def handleAdd (lhs rhs: Boogie.Expression.Expr) : Boogie.Expression.Expr :=
   | (.tcons "string" []), (.tcons "string" []) => .app () (.app () (.op () "Str.Concat" mty[string → (string → string)]) lhs) rhs
   | _, _ => panic! s!"Unimplemented add op for {lhs} + {rhs}"
 
-partial def PyExprToBoogie (e : Python.expr SourceRange) : Boogie.Expression.Expr :=
-  match e with
-  | .Call _ _ _ _ => panic! s!"Call should be handled at stmt level: {repr e}"
-  | .Constant _ c _ => PyConstToBoogie c
-  | .Name _ n _ =>
-    match n.val with
-    | "AssertionError" | "Exception" => .strConst () n.val
-    | _ => .fvar () n.val none
-  | .JoinedStr _ ss => PyExprToBoogie ss.val[0]! -- TODO: need to actually join strings
-  | .BinOp _ lhs op rhs => match op with
-    | .Add _ => handleAdd (PyExprToBoogie lhs) (PyExprToBoogie rhs)
-    | _ => panic! s!"Unhandled BinOp: {repr e}"
-  | .Compare _ lhs op rhs =>
-    match op.val with
-    | #[v] => match v with
-      | Strata.Python.cmpop.Eq _ =>
-        let l := PyExprToBoogie lhs
-        assert! rhs.val.size == 1
-        let r := PyExprToBoogie rhs.val[0]!
-        (.eq () l r)
+structure SubstitutionRecord where
+  pyExpr : Python.expr SourceRange
+  boogieExpr : Boogie.Expression.Expr
+
+instance : Repr (List SubstitutionRecord) where
+  reprPrec xs _ :=
+    let py_exprs := xs.map (λ r => r.pyExpr)
+    s!"{repr py_exprs}"
+
+def PyExprIdent (e1 e2: Python.expr SourceRange) : Bool :=
+  match e1, e2 with
+  | .Call sr1 _ _ _, .Call sr2 _ _ _ => sr1 == sr2
+  | _ , _ => false
+
+partial def PyExprToBoogie (e : Python.expr SourceRange) (substitution_records : Option (List SubstitutionRecord) := none) : Boogie.Expression.Expr :=
+  if h : substitution_records.isSome && (substitution_records.get!.find? (λ r => PyExprIdent r.pyExpr e)).isSome then
+    have hr : (List.find? (fun r => PyExprIdent r.pyExpr e) substitution_records.get!).isSome = true := by rw [Bool.and_eq_true] at h; exact h.2
+    let record := (substitution_records.get!.find? (λ r => PyExprIdent r.pyExpr e)).get hr
+    record.boogieExpr
+  else
+    match e with
+    | .Call _ f _ _ => panic! s!"Call should be handled at stmt level: \n(func: {repr f}) \n(Records: {repr substitution_records}) \n(AST: {repr e.toAst})"
+    | .Constant _ c _ => PyConstToBoogie c
+    | .Name _ n _ =>
+      match n.val with
+      | "AssertionError" | "Exception" => .strConst () n.val
+      | _ => .fvar () n.val none
+    | .JoinedStr _ ss => PyExprToBoogie ss.val[0]! -- TODO: need to actually join strings
+    | .BinOp _ lhs op rhs => match op with
+      | .Add _ => handleAdd (PyExprToBoogie lhs) (PyExprToBoogie rhs)
+      | _ => panic! s!"Unhandled BinOp: {repr e}"
+    | .Compare _ lhs op rhs =>
+      match op.val with
+      | #[v] => match v with
+        | Strata.Python.cmpop.Eq _ =>
+          let l := PyExprToBoogie lhs
+          assert! rhs.val.size == 1
+          let r := PyExprToBoogie rhs.val[0]!
+          (.eq () l r)
+        | _ => panic! s!"Unhandled comparison op: {repr op.val}"
       | _ => panic! s!"Unhandled comparison op: {repr op.val}"
-    | _ => panic! s!"Unhandled comparison op: {repr op.val}"
-  | _ => panic! s!"Unhandled Expr: {repr e}"
+    | _ => panic! s!"Unhandled Expr: {repr e}"
+
+partial def PyExprToBoogieWithSubst (substitution_records : Option (List SubstitutionRecord)) (e : Python.expr SourceRange) : Boogie.Expression.Expr :=
+  PyExprToBoogie e substitution_records
 
 partial def PyExprToString (e : Python.expr SourceRange) : String :=
   match e with
@@ -114,11 +136,11 @@ partial def PyExprToString (e : Python.expr SourceRange) : String :=
     | _ => panic! s!"Unsupported subscript to string: {repr e}"
   | _ => panic! s!"Unhandled Expr: {repr e}"
 
-partial def PyKWordsToBoogie (kw : Python.keyword SourceRange) : (String × Boogie.Expression.Expr) :=
+partial def PyKWordsToBoogie (substitution_records : Option (List SubstitutionRecord)) (kw : Python.keyword SourceRange) : (String × Boogie.Expression.Expr) :=
   match kw with
   | .mk_keyword _ name expr =>
     match name.val with
-    | some n => (n.val, PyExprToBoogie expr)
+    | some n => (n.val, PyExprToBoogieWithSubst substitution_records expr)
     | none => panic! "Keyword arg should have a name"
 
 structure PythonFunctionDecl where
@@ -139,17 +161,23 @@ def callCanThrow (func_infos : List PythonFunctionDecl) (stmt: Python.stmt Sourc
 
 -- TODO: we should be checking that args are right
 open Strata.Python.Internal in
-def argsAndKWordsToCanonicalList (func_infos : List PythonFunctionDecl) (fname: String) (args : Array (Python.expr SourceRange)) (kwords: Array (Python.keyword SourceRange)) : List Boogie.Expression.Expr :=
+def argsAndKWordsToCanonicalList (func_infos : List PythonFunctionDecl)
+                                 (fname: String)
+                                 (args : Array (Python.expr SourceRange))
+                                 (kwords: Array (Python.keyword SourceRange))
+                                 (substitution_records : Option (List SubstitutionRecord) := none) : List Boogie.Expression.Expr :=
+  dbg_trace s!"Canonicalizing {fname} (substitution_records: {substitution_records.get!.length})"
+
   -- TODO: we need a more general solution for other functions
   if fname == "print" then
     args.toList.map PyExprToBoogie
   else if func_infos.any (λ e => e.name == fname) then
-    args.toList.map PyExprToBoogie
+    args.toList.map (PyExprToBoogieWithSubst substitution_records)
   else
     let required_order := getFuncSigOrder fname
     assert! args.size <= required_order.length
     let remaining := required_order.drop args.size
-    let kws_and_exprs := kwords.toList.map PyKWordsToBoogie
+    let kws_and_exprs := kwords.toList.map (PyKWordsToBoogie substitution_records)
     let ordered_remaining_args := remaining.map (λ n => match kws_and_exprs.find? (λ p => p.fst == n) with
       | .some p =>
         let type_str := getFuncSigType fname n
@@ -162,7 +190,7 @@ def argsAndKWordsToCanonicalList (func_infos : List PythonFunctionDecl) (fname:
         else
           p.snd
       | .none => Strata.Python.TypeStrToBoogieExpr (getFuncSigType fname n))
-    args.toList.map PyExprToBoogie ++ ordered_remaining_args
+    args.toList.map (PyExprToBoogieWithSubst substitution_records) ++ ordered_remaining_args
 
 def handleCallThrow (jmp_target : String) : Boogie.Statement :=
   let cond := .eq () (.app () (.op () "ExceptOrNone_tag" none) (.fvar () "maybe_except" none)) (.op () "EN_STR_TAG" none)
@@ -219,6 +247,14 @@ def collectVarDecls (stmts: Array (Python.stmt SourceRange)) : List Boogie.State
   let foo := dedup.map toBoogie
   foo.flatten
 
+def isCall (e: Python.expr SourceRange) : Bool :=
+  match e with
+  | .Call _ _ _ _ => true
+  | _ => false
+
+def initTmpParam (p: Python.expr SourceRange × String) : Boogie.Statement :=
+  .init p.snd t[string] (PyExprToBoogie p.fst)
+
 mutual
 
 partial def exceptHandlersToBoogie (jmp_targets: List String) (func_infos : List PythonFunctionDecl) (h : Python.excepthandler SourceRange) : List Boogie.Statement :=
@@ -240,6 +276,27 @@ partial def exceptHandlersToBoogie (jmp_targets: List String) (func_infos : List
     let body_if_matches := body.val.toList.flatMap (PyStmtToBoogie jmp_targets func_infos) ++ [.goto jmp_targets[1]!]
     set_ex_ty_matches ++ [.ite cond {ss := body_if_matches} {ss := []}]
 
+partial def handleFunctionCall (lhs: List Boogie.Expression.Ident)
+                               (fname: String)
+                               (args: Ann (Array (Python.expr SourceRange)) SourceRange)
+                               (kwords: Ann (Array (Python.keyword SourceRange)) SourceRange)
+                               (_jmp_targets: List String)
+                               (func_infos : List PythonFunctionDecl)
+                               (_s : Python.stmt SourceRange) : List Boogie.Statement :=
+  -- Boogie doesn't allow nested function calls, so we need to introduce temporary variables for each nested call
+  let nested_args_calls := args.val.filterMap (λ a => if isCall a then some a else none)
+  let args_calls_to_tmps := nested_args_calls.map (λ a => (a, s!"call_arg_tmp_{a.toAst.ann.start}"))
+  let nested_kwords_calls := kwords.val.filterMap (λ a =>
+    let arg := match a with
+      | .mk_keyword _ _ f => f
+    if isCall arg then some arg else none)
+  let kwords_calls_to_tmps := nested_kwords_calls.map (λ a => (a, s!"call_kword_tmp_{a.toAst.ann.start}"))
+
+  let substitution_records : List SubstitutionRecord := args_calls_to_tmps.toList.map (λ p => {pyExpr := p.fst, boogieExpr := .fvar () p.snd none}) ++
+                                                        kwords_calls_to_tmps.toList.map (λ p => {pyExpr := p.fst, boogieExpr := .fvar () p.snd none})
+  args_calls_to_tmps.toList.map initTmpParam ++
+    kwords_calls_to_tmps.toList.map initTmpParam ++
+    [.call lhs fname (argsAndKWordsToCanonicalList func_infos fname args.val kwords.val substitution_records)]
 
 partial def PyStmtToBoogie (jmp_targets: List String) (func_infos : List PythonFunctionDecl) (s : Python.stmt SourceRange) : List Boogie.Statement :=
   assert! jmp_targets.length > 0
@@ -257,21 +314,21 @@ partial def PyStmtToBoogie (jmp_targets: List String) (func_infos : List PythonF
     | .Expr _ (.Call _ func args kwords) =>
       let fname := PyExprToString func
       if callCanThrow func_infos s then
-        [.call ["maybe_except"] fname (argsAndKWordsToCanonicalList func_infos fname args.val kwords.val)]
+        handleFunctionCall ["maybe_except"] fname args kwords jmp_targets func_infos s
       else
-        [.call [] fname (argsAndKWordsToCanonicalList func_infos fname args.val kwords.val)]
+        handleFunctionCall [] fname args kwords jmp_targets func_infos s
     | .Expr _ _ =>
       panic! "Can't handle Expr statements that aren't calls"
     | .Assign _ lhs (.Call _ func args kwords) _ =>
       assert! lhs.val.size == 1
       let fname := PyExprToString func
-      [.call [PyExprToString lhs.val[0]!, "maybe_except"] fname (argsAndKWordsToCanonicalList func_infos fname args.val kwords.val)]
+      handleFunctionCall [PyExprToString lhs.val[0]!, "maybe_except"] fname args kwords jmp_targets func_infos s
     | .Assign _ lhs rhs _ =>
       assert! lhs.val.size == 1
       [.set (PyExprToString lhs.val[0]!) (PyExprToBoogie rhs)]
     | .AnnAssign _ lhs _ { ann := _ , val := (.some (.Call _ func args kwords))} _ =>
       let fname := PyExprToString func
-      [.call [PyExprToString lhs, "maybe_except"] fname (argsAndKWordsToCanonicalList func_infos fname args.val kwords.val)]
+      handleFunctionCall [PyExprToString lhs, "maybe_except"] fname args kwords jmp_targets func_infos s
     | .AnnAssign _ lhs _ {ann := _, val := (.some e)} _ =>
       [.set (PyExprToString lhs) (PyExprToBoogie e)]
     | .Try _ body handlers _orelse _finalbody =>

From 74803b280992d1caec2752031bdb8387ba1a5f56 Mon Sep 17 00:00:00 2001
From: Andrew Wells <anmwells@amazon.com>
Date: Mon, 1 Dec 2025 14:32:40 -0600
Subject: [PATCH 3/8] Handle nested function calls

---
 Strata/Languages/Python/BoogiePrelude.lean    |  5 ++-
 .../Languages/Python/FunctionSignatures.lean  |  6 +++
 Strata/Languages/Python/PythonToBoogie.lean   | 40 ++++++++++++++-----
 3 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/Strata/Languages/Python/BoogiePrelude.lean b/Strata/Languages/Python/BoogiePrelude.lean
index 2a704e813..d5589b4be 100644
--- a/Strata/Languages/Python/BoogiePrelude.lean
+++ b/Strata/Languages/Python/BoogiePrelude.lean
@@ -142,7 +142,10 @@ procedure importFrom(module : string, names : ListStr, level : int) returns ()
 procedure import(names : ListStr) returns ()
 ;
 
-procedure print(msg : string) returns ()
+procedure print(msg : string, opt : StrOrNone) returns ()
+;
+
+procedure json_dumps(msg : string) returns (s: string)
 ;
 
 function str_len(s : string) : int;
diff --git a/Strata/Languages/Python/FunctionSignatures.lean b/Strata/Languages/Python/FunctionSignatures.lean
index 4a3e076a0..07ffec0fb 100644
--- a/Strata/Languages/Python/FunctionSignatures.lean
+++ b/Strata/Languages/Python/FunctionSignatures.lean
@@ -13,6 +13,7 @@ namespace Python
 def getFuncSigOrder (fname: String) : List String :=
   match fname with
   | "test_helper_procedure" => ["req_name", "opt_name"]
+  | "print" => ["msg", "opt"]
   | _ => panic! s!"Missing function signature : {fname}"
 
 -- We should extract the function signatures from the prelude:
@@ -23,6 +24,11 @@ def getFuncSigType (fname: String) (arg: String) : String :=
     | "req_name" => "string"
     | "opt_name" => "StrOrNone"
     | _ => panic! s!"Unrecognized arg : {arg}"
+  | "print" =>
+    match arg with
+    | "msg" => "string"
+    | "opt" => "StrOrNone"
+    | _ => panic! s!"Unrecognized arg : {arg}"
   | _ => panic! s!"Missing function signature : {fname}"
 
 def TypeStrToBoogieExpr (ty: String) : Boogie.Expression.Expr :=
diff --git a/Strata/Languages/Python/PythonToBoogie.lean b/Strata/Languages/Python/PythonToBoogie.lean
index 14043a654..2543e5ace 100644
--- a/Strata/Languages/Python/PythonToBoogie.lean
+++ b/Strata/Languages/Python/PythonToBoogie.lean
@@ -166,12 +166,14 @@ def argsAndKWordsToCanonicalList (func_infos : List PythonFunctionDecl)
                                  (args : Array (Python.expr SourceRange))
                                  (kwords: Array (Python.keyword SourceRange))
                                  (substitution_records : Option (List SubstitutionRecord) := none) : List Boogie.Expression.Expr :=
-  dbg_trace s!"Canonicalizing {fname} (substitution_records: {substitution_records.get!.length})"
-
   -- TODO: we need a more general solution for other functions
-  if fname == "print" then
-    args.toList.map PyExprToBoogie
-  else if func_infos.any (λ e => e.name == fname) then
+  -- if fname == "print" then
+  --   if args.size == 1 then
+  --     args.toList.map PyExprToBoogie
+  --   else
+  --     args.toList.map PyExprToBoogie
+  -- else if
+  if func_infos.any (λ e => e.name == fname) then
     args.toList.map (PyExprToBoogieWithSubst substitution_records)
   else
     let required_order := getFuncSigOrder fname
@@ -190,7 +192,22 @@ def argsAndKWordsToCanonicalList (func_infos : List PythonFunctionDecl)
         else
           p.snd
       | .none => Strata.Python.TypeStrToBoogieExpr (getFuncSigType fname n))
-    args.toList.map (PyExprToBoogieWithSubst substitution_records) ++ ordered_remaining_args
+    let args := args.map (PyExprToBoogieWithSubst substitution_records)
+    let args := (List.range required_order.length).filterMap (λ n =>
+        if h: n < args.size then
+          let arg_name := required_order[n]!
+          let type_str := getFuncSigType fname arg_name
+          if type_str.endsWith "OrNone" then
+            -- Optional param. Need to wrap e.g., string into StrOrNone
+            match type_str with
+            | "StrOrNone" => some (.app () (.op () "StrOrNone_mk_str" none) args[n]!)
+            | "BytesOrStrOrNone" => some (.app () (.op () "BytesOrStrOrNone_mk_str" none) args[n]!)
+            | _ => panic! "Unsupported type_str: "++ type_str
+          else
+            some (args[n]!)
+        else
+          none)
+    args ++ ordered_remaining_args
 
 def handleCallThrow (jmp_target : String) : Boogie.Statement :=
   let cond := .eq () (.app () (.op () "ExceptOrNone_tag" none) (.fvar () "maybe_except" none)) (.op () "EN_STR_TAG" none)
@@ -252,8 +269,11 @@ def isCall (e: Python.expr SourceRange) : Bool :=
   | .Call _ _ _ _ => true
   | _ => false
 
-def initTmpParam (p: Python.expr SourceRange × String) : Boogie.Statement :=
-  .init p.snd t[string] (PyExprToBoogie p.fst)
+def initTmpParam (p: Python.expr SourceRange × String) : List Boogie.Statement :=
+-- [.call lhs fname (argsAndKWordsToCanonicalList func_infos fname args.val kwords.val substitution_records)]
+  match p.fst with
+  | .Call _ f args _ => [(.init p.snd t[string] (.strConst () "trash")), .call [p.snd] "json_dumps" [(.strConst () "dummy")]]
+  | _ => panic! "Expected Call"
 
 mutual
 
@@ -294,8 +314,8 @@ partial def handleFunctionCall (lhs: List Boogie.Expression.Ident)
 
   let substitution_records : List SubstitutionRecord := args_calls_to_tmps.toList.map (λ p => {pyExpr := p.fst, boogieExpr := .fvar () p.snd none}) ++
                                                         kwords_calls_to_tmps.toList.map (λ p => {pyExpr := p.fst, boogieExpr := .fvar () p.snd none})
-  args_calls_to_tmps.toList.map initTmpParam ++
-    kwords_calls_to_tmps.toList.map initTmpParam ++
+  args_calls_to_tmps.toList.flatMap initTmpParam ++
+    kwords_calls_to_tmps.toList.flatMap initTmpParam ++
     [.call lhs fname (argsAndKWordsToCanonicalList func_infos fname args.val kwords.val substitution_records)]
 
 partial def PyStmtToBoogie (jmp_targets: List String) (func_infos : List PythonFunctionDecl) (s : Python.stmt SourceRange) : List Boogie.Statement :=

From 9085e2e18365cd7b6c1fc043a8b944d260f73b4b Mon Sep 17 00:00:00 2001
From: Andrew Wells <anmwells@amazon.com>
Date: Mon, 1 Dec 2025 14:34:13 -0600
Subject: [PATCH 4/8] Remove some unproved array bounds

---
 Strata/Languages/Python/PythonToBoogie.lean | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/Strata/Languages/Python/PythonToBoogie.lean b/Strata/Languages/Python/PythonToBoogie.lean
index 2543e5ace..c73b6493c 100644
--- a/Strata/Languages/Python/PythonToBoogie.lean
+++ b/Strata/Languages/Python/PythonToBoogie.lean
@@ -195,16 +195,16 @@ def argsAndKWordsToCanonicalList (func_infos : List PythonFunctionDecl)
     let args := args.map (PyExprToBoogieWithSubst substitution_records)
     let args := (List.range required_order.length).filterMap (λ n =>
         if h: n < args.size then
-          let arg_name := required_order[n]!
+          let arg_name := required_order[n]! -- Guaranteed by range. Using finRange causes breaking coercions to Nat.
           let type_str := getFuncSigType fname arg_name
           if type_str.endsWith "OrNone" then
             -- Optional param. Need to wrap e.g., string into StrOrNone
             match type_str with
-            | "StrOrNone" => some (.app () (.op () "StrOrNone_mk_str" none) args[n]!)
-            | "BytesOrStrOrNone" => some (.app () (.op () "BytesOrStrOrNone_mk_str" none) args[n]!)
+            | "StrOrNone" => some (.app () (.op () "StrOrNone_mk_str" none) args[n])
+            | "BytesOrStrOrNone" => some (.app () (.op () "BytesOrStrOrNone_mk_str" none) args[n])
             | _ => panic! "Unsupported type_str: "++ type_str
           else
-            some (args[n]!)
+            some (args[n])
         else
           none)
     args ++ ordered_remaining_args

From 1efd06c8d26eeaf93b65c22e639419dc5e6fc655 Mon Sep 17 00:00:00 2001
From: Andrew Wells <anmwells@amazon.com>
Date: Tue, 2 Dec 2025 13:50:32 -0600
Subject: [PATCH 5/8] Return values, dictionaries; json_loads, json_dumps and
 input; var decls in If bodies

---
 Strata/Languages/Python/BoogiePrelude.lean    |  9 ++++-
 .../Languages/Python/FunctionSignatures.lean  | 16 ++++++++
 Strata/Languages/Python/PythonToBoogie.lean   | 38 +++++++++++--------
 3 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/Strata/Languages/Python/BoogiePrelude.lean b/Strata/Languages/Python/BoogiePrelude.lean
index d5589b4be..433550f6a 100644
--- a/Strata/Languages/Python/BoogiePrelude.lean
+++ b/Strata/Languages/Python/BoogiePrelude.lean
@@ -104,6 +104,7 @@ function BoolOrStrOrNone_mk_bool(b : bool) : (BoolOrStrOrNone);
 function BoolOrStrOrNone_mk_str(s : string) : (BoolOrStrOrNone);
 function BoolOrStrOrNone_mk_none(v : None) : (BoolOrStrOrNone);
 function Client_tag(v : Client) : (ClientTag);
+function DictStrAny_mk(s : string) : (DictStrAny);
 
 // Unique const axioms
 axiom [unique_ExceptOrNoneTag]: EN_STR_TAG != EN_NONE_TAG;
@@ -145,7 +146,13 @@ procedure import(names : ListStr) returns ()
 procedure print(msg : string, opt : StrOrNone) returns ()
 ;
 
-procedure json_dumps(msg : string) returns (s: string)
+procedure json_dumps(msg : DictStrAny, opt_indent : IntOrNone) returns (s: string, maybe_except: ExceptOrNone)
+;
+
+procedure json_loads(msg : string) returns (d: DictStrAny, maybe_except: ExceptOrNone)
+;
+
+procedure input(msg : string) returns (result: string, maybe_except: ExceptOrNone)
 ;
 
 function str_len(s : string) : int;
diff --git a/Strata/Languages/Python/FunctionSignatures.lean b/Strata/Languages/Python/FunctionSignatures.lean
index 07ffec0fb..dfccf4d22 100644
--- a/Strata/Languages/Python/FunctionSignatures.lean
+++ b/Strata/Languages/Python/FunctionSignatures.lean
@@ -14,6 +14,9 @@ def getFuncSigOrder (fname: String) : List String :=
   match fname with
   | "test_helper_procedure" => ["req_name", "opt_name"]
   | "print" => ["msg", "opt"]
+  | "json_dumps" => ["msg", "opt_indent"]
+  | "json_loads" => ["msg"]
+  | "input" => ["msg"]
   | _ => panic! s!"Missing function signature : {fname}"
 
 -- We should extract the function signatures from the prelude:
@@ -29,6 +32,19 @@ def getFuncSigType (fname: String) (arg: String) : String :=
     | "msg" => "string"
     | "opt" => "StrOrNone"
     | _ => panic! s!"Unrecognized arg : {arg}"
+  | "json_dumps" =>
+    match arg with
+    | "msg" => "DictStrAny"
+    | "opt_indent" => "IntOrNone"
+    | _ => panic! s!"Unrecognized arg : {arg}"
+  | "json_loads" =>
+    match arg with
+    | "msg" => "string"
+    | _ => panic! s!"Unrecognized arg : {arg}"
+  | "input" =>
+    match arg with
+    | "msg" => "string"
+    | _ => panic! s!"Unrecognized arg : {arg}"
   | _ => panic! s!"Missing function signature : {fname}"
 
 def TypeStrToBoogieExpr (ty: String) : Boogie.Expression.Expr :=
diff --git a/Strata/Languages/Python/PythonToBoogie.lean b/Strata/Languages/Python/PythonToBoogie.lean
index c73b6493c..fb1a4b404 100644
--- a/Strata/Languages/Python/PythonToBoogie.lean
+++ b/Strata/Languages/Python/PythonToBoogie.lean
@@ -74,6 +74,9 @@ def handleAdd (lhs rhs: Boogie.Expression.Expr) : Boogie.Expression.Expr :=
   | (.tcons "string" []), (.tcons "string" []) => .app () (.app () (.op () "Str.Concat" mty[string → (string → string)]) lhs) rhs
   | _, _ => panic! s!"Unimplemented add op for {lhs} + {rhs}"
 
+def handleDict (keys: Array (Python.opt_expr SourceRange)) (values: Array (Python.expr SourceRange)) : Boogie.Expression.Expr :=
+  .app () (.op () "DictStrAny_mk" none) (.strConst () "DefaultDict")
+
 structure SubstitutionRecord where
   pyExpr : Python.expr SourceRange
   boogieExpr : Boogie.Expression.Expr
@@ -115,6 +118,7 @@ partial def PyExprToBoogie (e : Python.expr SourceRange) (substitution_records :
           (.eq () l r)
         | _ => panic! s!"Unhandled comparison op: {repr op.val}"
       | _ => panic! s!"Unhandled comparison op: {repr op.val}"
+    | .Dict _ keys values => handleDict keys.val values.val
     | _ => panic! s!"Unhandled Expr: {repr e}"
 
 partial def PyExprToBoogieWithSubst (substitution_records : Option (List SubstitutionRecord)) (e : Python.expr SourceRange) : Boogie.Expression.Expr :=
@@ -166,13 +170,6 @@ def argsAndKWordsToCanonicalList (func_infos : List PythonFunctionDecl)
                                  (args : Array (Python.expr SourceRange))
                                  (kwords: Array (Python.keyword SourceRange))
                                  (substitution_records : Option (List SubstitutionRecord) := none) : List Boogie.Expression.Expr :=
-  -- TODO: we need a more general solution for other functions
-  -- if fname == "print" then
-  --   if args.size == 1 then
-  --     args.toList.map PyExprToBoogie
-  --   else
-  --     args.toList.map PyExprToBoogie
-  -- else if
   if func_infos.any (λ e => e.name == fname) then
     args.toList.map (PyExprToBoogieWithSubst substitution_records)
   else
@@ -239,14 +236,15 @@ def deduplicateTypeAnnotations (l : List (String × Option String)) : List (Stri
     | .some ty => (n, ty)
     | .none => panic s!"Missing type annotations for {n}")
 
-def collectVarDecls (stmts: Array (Python.stmt SourceRange)) : List Boogie.Statement :=
-  let go (s : Python.stmt SourceRange) : List (String × Option String) :=
+partial def collectVarDecls (stmts: Array (Python.stmt SourceRange)) : List Boogie.Statement :=
+  let rec go (s : Python.stmt SourceRange) : List (String × Option String) :=
     match s with
     | .Assign _ lhs _ _ =>
       let names := lhs.val.toList.map PyExprToString
       names.map (λ n => (n, none))
     | .AnnAssign _ lhs ty _ _ =>
       [(PyExprToString lhs, PyExprToString ty)]
+    | .If _ _ body _ => body.val.toList.flatMap go
     | _ => []
   let dup := stmts.toList.flatMap go
   let dedup := deduplicateTypeAnnotations dup
@@ -272,7 +270,8 @@ def isCall (e: Python.expr SourceRange) : Bool :=
 def initTmpParam (p: Python.expr SourceRange × String) : List Boogie.Statement :=
 -- [.call lhs fname (argsAndKWordsToCanonicalList func_infos fname args.val kwords.val substitution_records)]
   match p.fst with
-  | .Call _ f args _ => [(.init p.snd t[string] (.strConst () "trash")), .call [p.snd] "json_dumps" [(.strConst () "dummy")]]
+  | .Call _ f args _ =>
+    [(.init p.snd t[string] (.strConst () "")), .call [p.snd, "maybe_except"] "json_dumps" [(.app () (.op () "DictStrAny_mk" none) (.strConst () "DefaultDict")), (Strata.Python.TypeStrToBoogieExpr "IntOrNone")]]
   | _ => panic! "Expected Call"
 
 mutual
@@ -361,6 +360,10 @@ partial def PyStmtToBoogie (jmp_targets: List String) (func_infos : List PythonF
     | .FunctionDef _ _ _ _ _ _ _ _ => panic! "Can't translate FunctionDef to Boogie statement"
     | .If _ test then_b else_b =>
       [.ite (PyExprToBoogie test) {ss := (ArrPyStmtToBoogie func_infos then_b.val)} {ss := (ArrPyStmtToBoogie func_infos else_b.val)}] -- TODO: fix this
+    | .Return _ v =>
+      match v.val with
+      | .some v => [.set "ret" (PyExprToBoogie v), .goto jmp_targets[0]!] -- TODO: need to thread return value name here. For now, assume "ret"
+      | .none => [.goto jmp_targets[0]!]
     | _ =>
       panic! s!"Unsupported {repr s}"
   if callCanThrow func_infos s then
@@ -397,16 +400,19 @@ def pyTyStrToLMonoTy (ty_str: String) : Lambda.LMonoTy :=
   | "str" => mty[string]
   | _ => panic! s!"Unsupported type: {ty_str}"
 
-def pythonFuncToBoogie (name : String) (args: List (String × String)) (body: Array (Python.stmt SourceRange)) (spec : Boogie.Procedure.Spec) (func_infos : List PythonFunctionDecl) : Boogie.Procedure :=
+def pythonFuncToBoogie (name : String) (args: List (String × String)) (body: Array (Python.stmt SourceRange)) (ret : Option (Python.expr SourceRange)) (spec : Boogie.Procedure.Spec) (func_infos : List PythonFunctionDecl) : Boogie.Procedure :=
   let inputs : List (Lambda.Identifier Boogie.Visibility × Lambda.LMonoTy) := args.map (λ p => (p.fst, pyTyStrToLMonoTy p.snd))
   let varDecls := collectVarDecls body ++ [(.init "exception_ty_matches" t[bool] (.boolConst () false)), (.havoc "exception_ty_matches")]
   let stmts := ArrPyStmtToBoogie func_infos body
   let body := varDecls ++ stmts ++ [.block "end" {ss := []}]
+  let outputs : Lambda.LMonoTySignature := match ret with
+  | .some v => [("ret", (.tcons "DictStrAny" [])), ("maybe_except", (.tcons "ExceptOrNone" []))]
+  | .none => [("maybe_except", (.tcons "ExceptOrNone" []))]
   {
     header := {name,
                typeArgs := [],
                inputs,
-               outputs := [("maybe_except", (.tcons "ExceptOrNone" []))]},
+               outputs},
     spec,
     body
   }
@@ -421,13 +427,13 @@ def unpackPyArguments (args: Python.arguments SourceRange) : List (String × Str
     | .mk_arg _ name oty _ =>
       match oty.val with
       | .some ty => (name.val, PyExprToString ty)
-      | _ => panic! s!"Missing type annotation on arg: {repr a}")
+      | _ => panic! s!"Missing type annotation on arg: {repr a} ({repr args})")
 
 def PyFuncDefToBoogie (s: Python.stmt SourceRange) (func_infos : List PythonFunctionDecl) : Boogie.Decl × PythonFunctionDecl :=
   match s with
-  | .FunctionDef _ name args body _ _ret _ _ =>
+  | .FunctionDef _ name args body _ ret _ _ =>
     let args := unpackPyArguments args
-    (.proc (pythonFuncToBoogie name.val args body.val default func_infos), {name := name.val, args})
+    (.proc (pythonFuncToBoogie name.val args body.val ret.val default func_infos), {name := name.val, args})
   | _ => panic! s!"Expected function def: {repr s}"
 
 def pythonToBoogie (pgm: Strata.Program): Boogie.Program :=
@@ -458,6 +464,6 @@ def pythonToBoogie (pgm: Strata.Program): Boogie.Program :=
   let func_defs := func_defs_and_infos.fst
   let func_infos := func_defs_and_infos.snd
 
-  {decls := globals ++ func_defs ++ [.proc (pythonFuncToBoogie "__main__" [] non_func_blocks default func_infos)]}
+  {decls := globals ++ func_defs ++ [.proc (pythonFuncToBoogie "__main__" [] non_func_blocks none default func_infos)]}
 
 end Strata

From 6e503486e0ce09b12756c521228d0319d0f3a741 Mon Sep 17 00:00:00 2001
From: Andrew Wells <anmwells@amazon.com>
Date: Tue, 2 Dec 2025 16:31:37 -0600
Subject: [PATCH 6/8] List[str], for loops (unrolled once), list_comp, in,
 subscript added as uninterpreted fns

---
 Strata/Languages/Python/BoogiePrelude.lean    | 14 ++++
 .../Languages/Python/FunctionSignatures.lean  |  5 ++
 Strata/Languages/Python/PythonToBoogie.lean   | 69 ++++++++++++++-----
 3 files changed, 70 insertions(+), 18 deletions(-)

diff --git a/Strata/Languages/Python/BoogiePrelude.lean b/Strata/Languages/Python/BoogiePrelude.lean
index 433550f6a..c6110b753 100644
--- a/Strata/Languages/Python/BoogiePrelude.lean
+++ b/Strata/Languages/Python/BoogiePrelude.lean
@@ -88,6 +88,7 @@ function AnyOrNone_none_val(v : AnyOrNone) : (None);
 function AnyOrNone_mk_str(s : string) : (AnyOrNone);
 function AnyOrNone_mk_none(v : None) : (AnyOrNone);
 function IntOrNone_mk_none(v : None) : (IntOrNone);
+function IntOrNone_mk_int(i : int) : (IntOrNone);
 function BytesOrStrOrNone_mk_none(v : None) : (BytesOrStrOrNone);
 function BytesOrStrOrNone_mk_str(s : string) : (BytesOrStrOrNone);
 function MappingStrStrOrNone_mk_none(v : None) : (MappingStrStrOrNone);
@@ -155,8 +156,21 @@ procedure json_loads(msg : string) returns (d: DictStrAny, maybe_except: ExceptO
 procedure input(msg : string) returns (result: string, maybe_except: ExceptOrNone)
 ;
 
+procedure random_choice(l : ListStr) returns (result: string, maybe_except: ExceptOrNone)
+;
+
 function str_len(s : string) : int;
 
+function str_in_list_str(s : string, l: ListStr) : bool;
+
+function str_in_dict_str_any(s : string, l: DictStrAny) : bool;
+
+function list_str_get(l : ListStr, i: int) : string;
+
+function dict_str_any_get(d : DictStrAny, k: string) : DictStrAny;
+
+function dict_str_any_length(d : DictStrAny) : int;
+
 procedure test_helper_procedure(req_name : string, opt_name : StrOrNone) returns (maybe_except: ExceptOrNone)
 spec {
   requires [req_name_is_foo]: req_name == "foo";
diff --git a/Strata/Languages/Python/FunctionSignatures.lean b/Strata/Languages/Python/FunctionSignatures.lean
index dfccf4d22..470ddc0ad 100644
--- a/Strata/Languages/Python/FunctionSignatures.lean
+++ b/Strata/Languages/Python/FunctionSignatures.lean
@@ -17,6 +17,7 @@ def getFuncSigOrder (fname: String) : List String :=
   | "json_dumps" => ["msg", "opt_indent"]
   | "json_loads" => ["msg"]
   | "input" => ["msg"]
+  | "random_choice" => ["l"]
   | _ => panic! s!"Missing function signature : {fname}"
 
 -- We should extract the function signatures from the prelude:
@@ -45,6 +46,10 @@ def getFuncSigType (fname: String) (arg: String) : String :=
     match arg with
     | "msg" => "string"
     | _ => panic! s!"Unrecognized arg : {arg}"
+  | "random_choice" =>
+    match arg with
+    | "l" => "ListStr"
+    | _ => panic! s!"Unrecognized arg : {arg}"
   | _ => panic! s!"Missing function signature : {fname}"
 
 def TypeStrToBoogieExpr (ty: String) : Boogie.Expression.Expr :=
diff --git a/Strata/Languages/Python/PythonToBoogie.lean b/Strata/Languages/Python/PythonToBoogie.lean
index fb1a4b404..0de70d422 100644
--- a/Strata/Languages/Python/PythonToBoogie.lean
+++ b/Strata/Languages/Python/PythonToBoogie.lean
@@ -28,6 +28,9 @@ def dummyDictStrAny : Boogie.Expression.Expr := .fvar () "DUMMY_DICT_STR_ANY" no
 def strType : Boogie.Expression.Ty := .forAll [] (.tcons "string" [])
 def dummyStr : Boogie.Expression.Expr := .fvar () "DUMMY_STR" none
 
+def listStrType : Boogie.Expression.Ty := .forAll [] (.tcons "ListStr" [])
+def dummyListStr : Boogie.Expression.Expr := .fvar () "DUMMY_LIST_STR" none
+
 
 -------------------------------------------------------------------------------
 
@@ -74,6 +77,12 @@ def handleAdd (lhs rhs: Boogie.Expression.Expr) : Boogie.Expression.Expr :=
   | (.tcons "string" []), (.tcons "string" []) => .app () (.app () (.op () "Str.Concat" mty[string → (string → string)]) lhs) rhs
   | _, _ => panic! s!"Unimplemented add op for {lhs} + {rhs}"
 
+def handleNot (arg: Boogie.Expression.Expr) : Boogie.Expression.Expr :=
+  let ty : Lambda.LMonoTy := (.tcons "ListStr" [])
+  match ty with
+  | (.tcons "ListStr" []) => .eq () arg (.op () "ListStr_nil" none)
+  | _ => panic! s!"Unimplemented not op for {arg}"
+
 def handleDict (keys: Array (Python.opt_expr SourceRange)) (values: Array (Python.expr SourceRange)) : Boogie.Expression.Expr :=
   .app () (.op () "DictStrAny_mk" none) (.strConst () "DefaultDict")
 
@@ -116,9 +125,22 @@ partial def PyExprToBoogie (e : Python.expr SourceRange) (substitution_records :
           assert! rhs.val.size == 1
           let r := PyExprToBoogie rhs.val[0]!
           (.eq () l r)
+        | Strata.Python.cmpop.In _ =>
+          let l := PyExprToBoogie lhs
+          assert! rhs.val.size == 1
+          let r := PyExprToBoogie rhs.val[0]!
+          .app () (.app () (.op () "str_in_dict_str_any" none) l) r
         | _ => panic! s!"Unhandled comparison op: {repr op.val}"
       | _ => panic! s!"Unhandled comparison op: {repr op.val}"
     | .Dict _ keys values => handleDict keys.val values.val
+    | .ListComp _ keys values => panic! "ListComp must be handled at stmt level"
+    | .UnaryOp _ op arg => match op with
+      | .Not _ => handleNot (PyExprToBoogie arg)
+      | _ => panic! "Unsupported UnaryOp: {repr e}"
+    | .Subscript _ v slice _ =>
+      let l := PyExprToBoogie v
+      let k := PyExprToBoogie slice
+      .app () (.app () (.op () "dict_str_any_get" none) l) k
     | _ => panic! s!"Unhandled Expr: {repr e}"
 
 partial def PyExprToBoogieWithSubst (substitution_records : Option (List SubstitutionRecord)) (e : Python.expr SourceRange) : Boogie.Expression.Expr :=
@@ -137,6 +159,10 @@ partial def PyExprToString (e : Python.expr SourceRange) : String :=
         assert! elts.val.size == 2
         s!"Dict[{PyExprToString elts.val[0]!} {PyExprToString elts.val[1]!}]"
       | _ => panic! s!"Unsupported slice: {repr slice}"
+    | "List" =>
+      match slice with
+      | .Name _ id _ => s!"List[{id.val}]"
+      | _ => panic! s!"Unsupported slice: {repr slice}"
     | _ => panic! s!"Unsupported subscript to string: {repr e}"
   | _ => panic! s!"Unhandled Expr: {repr e}"
 
@@ -163,6 +189,19 @@ def callCanThrow (func_infos : List PythonFunctionDecl) (stmt: Python.stmt Sourc
     | _ => false
   | _ => false
 
+open Strata.Python.Internal in
+def noneOrExpr (fname n : String) (e: Boogie.Expression.Expr) : Boogie.Expression.Expr :=
+  let type_str := getFuncSigType fname n
+  if type_str.endsWith "OrNone" then
+    -- Optional param. Need to wrap e.g., string into StrOrNone
+    match type_str with
+    | "IntOrNone" => .app () (.op () "IntOrNone_mk_int" none) e
+    | "StrOrNone" => .app () (.op () "StrOrNone_mk_str" none) e
+    | "BytesOrStrOrNone" => .app () (.op () "BytesOrStrOrNone_mk_str" none) e
+    | _ => panic! "Unsupported type_str: "++ type_str
+  else
+    e
+
 -- TODO: we should be checking that args are right
 open Strata.Python.Internal in
 def argsAndKWordsToCanonicalList (func_infos : List PythonFunctionDecl)
@@ -179,29 +218,13 @@ def argsAndKWordsToCanonicalList (func_infos : List PythonFunctionDecl)
     let kws_and_exprs := kwords.toList.map (PyKWordsToBoogie substitution_records)
     let ordered_remaining_args := remaining.map (λ n => match kws_and_exprs.find? (λ p => p.fst == n) with
       | .some p =>
-        let type_str := getFuncSigType fname n
-        if type_str.endsWith "OrNone" then
-          -- Optional param. Need to wrap e.g., string into StrOrNone
-          match type_str with
-          | "StrOrNone" => .app () (.op () "StrOrNone_mk_str" none) p.snd
-          | "BytesOrStrOrNone" => .app () (.op () "BytesOrStrOrNone_mk_str" none) p.snd
-          | _ => panic! "Unsupported type_str: "++ type_str
-        else
-          p.snd
+        noneOrExpr fname n p.snd
       | .none => Strata.Python.TypeStrToBoogieExpr (getFuncSigType fname n))
     let args := args.map (PyExprToBoogieWithSubst substitution_records)
     let args := (List.range required_order.length).filterMap (λ n =>
         if h: n < args.size then
           let arg_name := required_order[n]! -- Guaranteed by range. Using finRange causes breaking coercions to Nat.
-          let type_str := getFuncSigType fname arg_name
-          if type_str.endsWith "OrNone" then
-            -- Optional param. Need to wrap e.g., string into StrOrNone
-            match type_str with
-            | "StrOrNone" => some (.app () (.op () "StrOrNone_mk_str" none) args[n])
-            | "BytesOrStrOrNone" => some (.app () (.op () "BytesOrStrOrNone_mk_str" none) args[n])
-            | _ => panic! "Unsupported type_str: "++ type_str
-          else
-            some (args[n])
+          some (noneOrExpr fname arg_name args[n]!)
         else
           none)
     args ++ ordered_remaining_args
@@ -258,6 +281,7 @@ partial def collectVarDecls (stmts: Array (Python.stmt SourceRange)) : List Boog
     | "bytes" => [(.init name t[string] (.strConst () "")), (.havoc name)]
     | "Client" => [(.init name clientType dummyClient), (.havoc name)]
     | "Dict[str Any]" => [(.init name dictStrAnyType dummyDictStrAny), (.havoc name)]
+    | "List[str]" => [(.init name listStrType dummyListStr), (.havoc name)]
     | _ => panic! s!"Unsupported type annotation: `{ty_name}`"
   let foo := dedup.map toBoogie
   foo.flatten
@@ -348,6 +372,9 @@ partial def PyStmtToBoogie (jmp_targets: List String) (func_infos : List PythonF
     | .AnnAssign _ lhs _ { ann := _ , val := (.some (.Call _ func args kwords))} _ =>
       let fname := PyExprToString func
       handleFunctionCall [PyExprToString lhs, "maybe_except"] fname args kwords jmp_targets func_infos s
+    | .AnnAssign _ lhs _ { ann := _ , val := (.some (.ListComp _ _ _))} _ =>
+      -- TODO: check for errors first
+      [.havoc (PyExprToString lhs)]
     | .AnnAssign _ lhs _ {ann := _, val := (.some e)} _ =>
       [.set (PyExprToString lhs) (PyExprToBoogie e)]
     | .Try _ body handlers _orelse _finalbody =>
@@ -364,6 +391,12 @@ partial def PyStmtToBoogie (jmp_targets: List String) (func_infos : List PythonF
       match v.val with
       | .some v => [.set "ret" (PyExprToBoogie v), .goto jmp_targets[0]!] -- TODO: need to thread return value name here. For now, assume "ret"
       | .none => [.goto jmp_targets[0]!]
+    | .For _ _tgt itr body _ _ =>
+      -- Do one unrolling:
+      let guard := .app () (.op () "Bool.Not" none) (.eq () (.app () (.op () "dict_str_any_length" none) (PyExprToBoogie itr)) (.intConst () 0))
+      -- let guard := .boolConst () true
+      [.ite guard {ss := (ArrPyStmtToBoogie func_infos body.val)} {ss := []}]
+      -- TODO: missing havoc
     | _ =>
       panic! s!"Unsupported {repr s}"
   if callCanThrow func_infos s then

From e7f861202d5fdc6465d3c55ce57559a5b9c24f75 Mon Sep 17 00:00:00 2001
From: Joe Hendrix <joehx@amazon.com>
Date: Tue, 2 Dec 2025 16:48:56 -0800
Subject: [PATCH 7/8] Update pyTranslate and pyAnalyze to use compile time
 Python dialect

---
 Strata/DDM/Ion.lean | 15 +++++++++++++++
 StrataMain.lean     | 26 +++++++++++++-------------
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/Strata/DDM/Ion.lean b/Strata/DDM/Ion.lean
index 45a9842ad..67e9ff1d6 100644
--- a/Strata/DDM/Ion.lean
+++ b/Strata/DDM/Ion.lean
@@ -1299,4 +1299,19 @@ def fromIonFragment (f : Ion.Fragment) (dialects : DialectMap) (dialect : Dialec
     commands := commands
   }
 
+def fromIon (dialects : DialectMap) (dialect : DialectName) (bytes : ByteArray) : Except String Strata.Program := do
+  let (hdr, frag) ←
+    match Strata.Ion.Header.parse bytes with
+    | .error msg =>
+      throw msg
+    | .ok p =>
+      pure p
+  match hdr with
+  | .dialect _ =>
+    throw "Expected a Strata program instead of a dialect."
+  | .program name => do
+    if name != dialect then
+      throw s!"{name} program found when {dialect} expected."
+    fromIonFragment frag dialects dialect
+
 end Program
diff --git a/StrataMain.lean b/StrataMain.lean
index 95b29f9de..28b30d303 100644
--- a/StrataMain.lean
+++ b/StrataMain.lean
@@ -159,16 +159,20 @@ def diffCommand : Command where
     | _, _ =>
       exitFailure "Cannot compare dialect def with another dialect/program."
 
+def readPythonStrata (path : String) : IO Strata.Program := do
+  let bytes ← Strata.Util.readBinInputSource path
+  if ! bytes.startsWith Ion.binaryVersionMarker then
+    exitFailure s!"pyAnalyze expected Ion file"
+  match Strata.Program.fromIon Strata.Python.Python_map Strata.Python.Python.name bytes with
+  | .ok p => pure p
+  | .error msg => exitFailure msg
+
 def pyTranslateCommand : Command where
   name := "pyTranslate"
   args := [ "file" ]
-  help := "Tranlate a Strata Python Ion file to Strata.Boogie. Write results to stdout."
-  callback := fun searchPath v => do
-    let (ld, pd) ← readFile searchPath v[0]
-    match pd with
-    | .dialect d =>
-      IO.print <| d.format ld.dialects
-    | .program pgm =>
+  help := "Translate a Strata Python Ion file to Strata.Boogie. Write results to stdout."
+  callback := fun _ v => do
+    let pgm ← readPythonStrata v[0]
     let preludePgm := Strata.Python.Internal.Boogie.prelude
     let bpgm := Strata.pythonToBoogie pgm
     let newPgm : Boogie.Program := { decls := preludePgm.decls ++ bpgm.decls }
@@ -178,13 +182,9 @@ def pyAnalyzeCommand : Command where
   name := "pyAnalyze"
   args := [ "file", "verbose" ]
   help := "Analyze a Strata Python Ion file. Write results to stdout."
-  callback := fun searchPath v => do
+  callback := fun _ v => do
     let verbose := v[1] == "1"
-    let (ld, pd) ← readFile searchPath v[0]
-    match pd with
-    | .dialect d =>
-      IO.print <| d.format ld.dialects
-    | .program pgm =>
+    let pgm ← readPythonStrata v[0]
     if verbose then
       IO.print pgm
     let preludePgm := Strata.Python.Internal.Boogie.prelude

From e4723b66ff9b76d349fbdd4fbfa949f87ed3ee49 Mon Sep 17 00:00:00 2001
From: Joe Hendrix <joehx@amazon.com>
Date: Thu, 4 Dec 2025 14:27:04 -0800
Subject: [PATCH 8/8] Update statement level translation to use monad with
 error so we can replace panics with non-fatal error reporting.

This also replaces function list structures with more efficient maps and makes the builtin signatures defined in a datatype.
---
 .../Languages/Python/FunctionSignatures.lean  | 138 ++++++---
 Strata/Languages/Python/PythonDialect.lean    |   9 +
 Strata/Languages/Python/PythonToBoogie.lean   | 274 +++++++++++-------
 StrataMain.lean                               |  48 ++-
 4 files changed, 314 insertions(+), 155 deletions(-)

diff --git a/Strata/Languages/Python/FunctionSignatures.lean b/Strata/Languages/Python/FunctionSignatures.lean
index 470ddc0ad..b42025fe5 100644
--- a/Strata/Languages/Python/FunctionSignatures.lean
+++ b/Strata/Languages/Python/FunctionSignatures.lean
@@ -9,48 +9,40 @@ import Strata.Languages.Boogie.Boogie
 namespace Strata
 namespace Python
 
--- We should extract the function signatures from the prelude:
-def getFuncSigOrder (fname: String) : List String :=
-  match fname with
-  | "test_helper_procedure" => ["req_name", "opt_name"]
-  | "print" => ["msg", "opt"]
-  | "json_dumps" => ["msg", "opt_indent"]
-  | "json_loads" => ["msg"]
-  | "input" => ["msg"]
-  | "random_choice" => ["l"]
-  | _ => panic! s!"Missing function signature : {fname}"
-
--- We should extract the function signatures from the prelude:
-def getFuncSigType (fname: String) (arg: String) : String :=
-  match fname with
-  | "test_helper_procedure" =>
-    match arg with
-    | "req_name" => "string"
-    | "opt_name" => "StrOrNone"
-    | _ => panic! s!"Unrecognized arg : {arg}"
-  | "print" =>
-    match arg with
-    | "msg" => "string"
-    | "opt" => "StrOrNone"
-    | _ => panic! s!"Unrecognized arg : {arg}"
-  | "json_dumps" =>
-    match arg with
-    | "msg" => "DictStrAny"
-    | "opt_indent" => "IntOrNone"
-    | _ => panic! s!"Unrecognized arg : {arg}"
-  | "json_loads" =>
-    match arg with
-    | "msg" => "string"
-    | _ => panic! s!"Unrecognized arg : {arg}"
-  | "input" =>
-    match arg with
-    | "msg" => "string"
-    | _ => panic! s!"Unrecognized arg : {arg}"
-  | "random_choice" =>
-    match arg with
-    | "l" => "ListStr"
-    | _ => panic! s!"Unrecognized arg : {arg}"
-  | _ => panic! s!"Missing function signature : {fname}"
+inductive PyType where
+| AnyOrNone
+| BoolOrNone
+| BoolOrStrOrNone
+| BytesOrStrOrNone
+| DictStrAny
+| IntOrNone
+| ListStr
+| MappingStrStrOrNone
+| StrOrNone
+| Str
+deriving Inhabited
+
+namespace PyType
+
+def toString (atp : PyType) : String :=
+  match atp with
+  | AnyOrNone => "AnyOrNone"
+  | BoolOrNone => "BoolOrNone"
+  | BoolOrStrOrNone => "BoolOrStrOrNone"
+  | BytesOrStrOrNone => "BytesOrStrOrNone"
+  | DictStrAny => "DictStrAny"
+  | IntOrNone => "IntOrNone"
+  | ListStr => "ListStr"
+  | MappingStrStrOrNone => "MappingStrStrOrNone"
+  | StrOrNone => "StrOrNone"
+  | Str => "string"
+
+def allowNone (atp : PyType) : Bool :=
+  match atp with
+  | Str => False
+  | _ => True
+
+end PyType
 
 def TypeStrToBoogieExpr (ty: String) : Boogie.Expression.Expr :=
   if !ty.endsWith "OrNone" then
@@ -66,5 +58,67 @@ def TypeStrToBoogieExpr (ty: String) : Boogie.Expression.Expr :=
     | "MappingStrStrOrNone" => .app () (.op () "MappingStrStrOrNone_mk_none" none) (.op () "None_none" none)
     | _ => panic! s!"unsupported type: {ty}"
 
+
+structure FuncSignature where
+  args : Array (String × PyType)
+  required_arg_count : Nat := args.size
+  indexMap : Std.HashMap String Nat
+deriving Inhabited
+
+namespace FuncSignature
+
+def ofArray (args : Array (String × PyType)) : FuncSignature where
+  args := args
+  indexMap := args.foldl (init := {}) fun m (name, _) => m.insert name m.size
+
+end FuncSignature
+
+/--
+Global database of function signatures
+-/
+structure FuncSigDB where
+  -- TODO: Migrate from a map from String to a map that incorporates Python module
+  -- information.
+  map : Std.HashMap String FuncSignature := {}
+deriving Inhabited
+
+namespace FuncSigDB
+
+instance : GetElem? FuncSigDB String FuncSignature (fun db name => name ∈ db.map) where
+  getElem db nm p := db.map[nm]'p
+  getElem? db nm := db.map[nm]?
+  getElem! db nm := db.map[nm]!
+
+def add! (db : FuncSigDB) (name : String) (args : List (String × PyType)) :=
+    assert! name ∉ db.map
+    { db with map := db.map.insert name (.ofArray args.toArray) }
+
+end FuncSigDB
+
+def funcSigs (db : FuncSigDB := {}) : FuncSigDB := Prod.snd <| Id.run <| StateT.run (s := db) do
+  let arg (name : String) (tp : PyType) : String × PyType := (name, tp)
+  let add name args := modify (·.add! name args)
+  add "input" [
+    arg "msg" .Str,
+  ]
+  add "json_dumps" [
+    arg "msg" .DictStrAny,
+    arg "opt_indent" .IntOrNone,
+  ]
+  add "json_loads" [
+    arg "msg" .Str
+  ]
+  add "print" [
+    arg "msg" .Str,
+    arg "opt" .StrOrNone,
+  ]
+  add "random_choice" [
+    arg "l" .ListStr
+  ]
+  add "test_helper_procedure" [
+    arg "req_name" .Str,
+    arg "opt_name" .StrOrNone,
+  ]
+
 end Python
 end Strata
diff --git a/Strata/Languages/Python/PythonDialect.lean b/Strata/Languages/Python/PythonDialect.lean
index a6be1c386..444026f7d 100644
--- a/Strata/Languages/Python/PythonDialect.lean
+++ b/Strata/Languages/Python/PythonDialect.lean
@@ -16,6 +16,15 @@ namespace Strata
 namespace Python
 #load_dialect "../../../Tools/Python/test_results/dialects/Python.dialect.st.ion"
 #strata_gen Python
+
+/--
+Return the annotation associated with an expression
+
+FIXME. #strata_gen should eventually generate a more optimized version of this automatically.
+-/
+def expr.ann {α} [Inhabited α] (expr : Python.expr α) : α := expr.toAst |>.ann
+
+
 end Python
 
 
diff --git a/Strata/Languages/Python/PythonToBoogie.lean b/Strata/Languages/Python/PythonToBoogie.lean
index 0de70d422..29135aaf7 100644
--- a/Strata/Languages/Python/PythonToBoogie.lean
+++ b/Strata/Languages/Python/PythonToBoogie.lean
@@ -16,6 +16,17 @@ import Strata.Languages.Python.Regex.ReToBoogie
 import StrataTest.Internal.InternalFunctionSignatures
 
 namespace Strata
+
+namespace Python
+
+structure ErrorMsg where
+  loc : SourceRange
+  message : String
+deriving Inhabited
+
+def throwPyError {m α} [MonadExcept ErrorMsg m] (loc : SourceRange) (msg : String) : m α :=
+  throw { loc := loc, message := msg }
+
 open Lambda.LTy.Syntax
 -- Some hard-coded things we'll need to fix later:
 
@@ -83,7 +94,7 @@ def handleNot (arg: Boogie.Expression.Expr) : Boogie.Expression.Expr :=
   | (.tcons "ListStr" []) => .eq () arg (.op () "ListStr_nil" none)
   | _ => panic! s!"Unimplemented not op for {arg}"
 
-def handleDict (keys: Array (Python.opt_expr SourceRange)) (values: Array (Python.expr SourceRange)) : Boogie.Expression.Expr :=
+def handleDict (_keys: Array (Python.opt_expr SourceRange)) (_values: Array (Python.expr SourceRange)) : Boogie.Expression.Expr :=
   .app () (.op () "DictStrAny_mk" none) (.strConst () "DefaultDict")
 
 structure SubstitutionRecord where
@@ -178,56 +189,95 @@ structure PythonFunctionDecl where
   args : List (String × String) -- Elements are (arg_name, arg_ty) where `arg_ty` is the string representation of the type in Python
 deriving Repr, BEq, Inhabited
 
+/--
+Structure that stores information required to translate PythonAST expresions/stamts into
+Boogie.
+-/
+structure PyTransContext where
+  sig_db : Python.FuncSigDB
+  func_info : Std.HashMap String PythonFunctionDecl := {}
+deriving Inhabited
+
+namespace PyTransContext
+
+protected def add! (ctx : PyTransContext) (d : PythonFunctionDecl) : PyTransContext :=
+  assert! d.name ∉ ctx.func_info
+  { ctx with func_info := ctx.func_info.insert d.name d }
+
+end PyTransContext
+
 -- This information should come from our prelude. For now, we use the fact that
 -- these functions are exactly the ones
 -- represented as `Call(Attribute(Name(...)))` in the AST (instead of `Call(Name(...))`).
-def callCanThrow (func_infos : List PythonFunctionDecl) (stmt: Python.stmt SourceRange) : Bool :=
+def callCanThrow (ctx : PyTransContext) (stmt: Python.stmt SourceRange) : Bool :=
   match stmt with
   | .Expr _ (.Call _ (.Attribute _ _ _ _) _ _) | .Assign _ _ (.Call _ (.Attribute _ _ _ _) _ _) _ => true
   | .Expr _ (.Call _ f _ _) | .Assign _ _ (.Call _ f _ _) _ => match f with
-    | .Name _ f _ => func_infos.any (λ fi => fi.name == f.val)
+    | .Name _ f _ =>
+      match ctx.func_info[f.val]? with
+      | some _ => true
+      | none => false
     | _ => false
   | _ => false
 
+abbrev PyTransM := ReaderT PyTransContext (StateT (Array Boogie.Statement) (Except ErrorMsg))
+
+def add_stmt (stmt : Boogie.Statement) : PyTransM Unit := do
+  modify (·.push stmt)
+
+def liftPyTrans (m : PyTransM Unit) : PyTransM (List Boogie.Statement) := do
+  match m |>.run (←read) |>.run #[] with
+  | .ok (_, a) => pure <| a.toList
+  | .error msg => throw msg
+
+def runPyTrans (ctx : PyTransContext) (m : PyTransM Unit) : Except ErrorMsg (List Boogie.Statement) :=
+  match m |>.run ctx |>.run #[] with
+  | .ok (_, a) => pure a.toList
+  | .error msg => throw msg
+
 open Strata.Python.Internal in
-def noneOrExpr (fname n : String) (e: Boogie.Expression.Expr) : Boogie.Expression.Expr :=
-  let type_str := getFuncSigType fname n
-  if type_str.endsWith "OrNone" then
+def noneOrExpr (e: Boogie.Expression.Expr) (tp : PyType) : Boogie.Expression.Expr :=
+  if tp.allowNone then
     -- Optional param. Need to wrap e.g., string into StrOrNone
-    match type_str with
-    | "IntOrNone" => .app () (.op () "IntOrNone_mk_int" none) e
-    | "StrOrNone" => .app () (.op () "StrOrNone_mk_str" none) e
-    | "BytesOrStrOrNone" => .app () (.op () "BytesOrStrOrNone_mk_str" none) e
-    | _ => panic! "Unsupported type_str: "++ type_str
+    match tp  with
+    | .IntOrNone => .app () (.op () "IntOrNone_mk_int" none) e
+    | .StrOrNone => .app () (.op () "StrOrNone_mk_str" none) e
+    | .BytesOrStrOrNone => .app () (.op () "BytesOrStrOrNone_mk_str" none) e
+    | _ => panic! "Unsupported type_str: "++ tp.toString
   else
     e
 
 -- TODO: we should be checking that args are right
 open Strata.Python.Internal in
-def argsAndKWordsToCanonicalList (func_infos : List PythonFunctionDecl)
-                                 (fname: String)
+def argsAndKWordsToCanonicalList (ctx : PyTransContext)
+                                 (func : Python.expr SourceRange)
                                  (args : Array (Python.expr SourceRange))
                                  (kwords: Array (Python.keyword SourceRange))
-                                 (substitution_records : Option (List SubstitutionRecord) := none) : List Boogie.Expression.Expr :=
-  if func_infos.any (λ e => e.name == fname) then
-    args.toList.map (PyExprToBoogieWithSubst substitution_records)
-  else
-    let required_order := getFuncSigOrder fname
-    assert! args.size <= required_order.length
-    let remaining := required_order.drop args.size
+                                 (substitution_records : Option (List SubstitutionRecord) := none)
+                                : Except ErrorMsg (String × List Boogie.Expression.Expr) :=
+  -- TODO: we need a more general solution for other functions
+  let fname := PyExprToString func
+  if fname ∈ ctx.func_info then
+    .ok (fname, args.toList.map (PyExprToBoogieWithSubst substitution_records))
+  else do
+    let some fsig := ctx.sig_db[fname]?
+      | throwPyError func.ann s!"Missing function signature : {fname}"
+    assert! args.size <= fsig.args.size
+    let remaining := fsig.args.drop args.size
     let kws_and_exprs := kwords.toList.map (PyKWordsToBoogie substitution_records)
-    let ordered_remaining_args := remaining.map (λ n => match kws_and_exprs.find? (λ p => p.fst == n) with
+    let ordered_remaining_args := remaining.map (λ (n, tp) =>
+      match kws_and_exprs.find? (λ(l, _) => l == n) with
       | .some p =>
-        noneOrExpr fname n p.snd
-      | .none => Strata.Python.TypeStrToBoogieExpr (getFuncSigType fname n))
+        noneOrExpr p.snd tp
+      | .none => Strata.Python.TypeStrToBoogieExpr tp.toString)
     let args := args.map (PyExprToBoogieWithSubst substitution_records)
-    let args := (List.range required_order.length).filterMap (λ n =>
+    let args := (List.range fsig.args.size).filterMap (λ n =>
         if h: n < args.size then
-          let arg_name := required_order[n]! -- Guaranteed by range. Using finRange causes breaking coercions to Nat.
-          some (noneOrExpr fname arg_name args[n]!)
+          let (_, arg_type) := fsig.args[n]! -- Guaranteed by range. Using finRange causes breaking coercions to Nat.
+          some (noneOrExpr (args[n]'h) arg_type)
         else
           none)
-    args ++ ordered_remaining_args
+    pure <| (fname, args ++ ordered_remaining_args.toList)
 
 def handleCallThrow (jmp_target : String) : Boogie.Statement :=
   let cond := .eq () (.app () (.op () "ExceptOrNone_tag" none) (.fvar () "maybe_except" none)) (.op () "EN_STR_TAG" none)
@@ -294,38 +344,38 @@ def isCall (e: Python.expr SourceRange) : Bool :=
 def initTmpParam (p: Python.expr SourceRange × String) : List Boogie.Statement :=
 -- [.call lhs fname (argsAndKWordsToCanonicalList func_infos fname args.val kwords.val substitution_records)]
   match p.fst with
-  | .Call _ f args _ =>
+  | .Call _ _f _args _ =>
     [(.init p.snd t[string] (.strConst () "")), .call [p.snd, "maybe_except"] "json_dumps" [(.app () (.op () "DictStrAny_mk" none) (.strConst () "DefaultDict")), (Strata.Python.TypeStrToBoogieExpr "IntOrNone")]]
   | _ => panic! "Expected Call"
 
 mutual
 
-partial def exceptHandlersToBoogie (jmp_targets: List String) (func_infos : List PythonFunctionDecl) (h : Python.excepthandler SourceRange) : List Boogie.Statement :=
+partial def exceptHandlersToBoogie (jmp_targets: List String) (h : Python.excepthandler SourceRange) : PyTransM Unit :=
   assert! jmp_targets.length >= 2
   match h with
-  | .ExceptHandler _ ex_ty _ body =>
-    let set_ex_ty_matches := match ex_ty.val with
+  | .ExceptHandler _ ex_ty _ body => do
+    match ex_ty.val with
     | .some ex_ty =>
       let inherits_from : Boogie.BoogieIdent := "inheritsFrom"
       let get_ex_tag : Boogie.BoogieIdent := "ExceptOrNone_code_val"
       let exception_ty : Boogie.Expression.Expr := .app () (.op () get_ex_tag none) (.fvar () "maybe_except" none)
       let rhs_curried : Boogie.Expression.Expr := .app () (.op () inherits_from none) exception_ty
       let rhs : Boogie.Expression.Expr := .app () rhs_curried ((PyExprToBoogie ex_ty))
-      let call := .set "exception_ty_matches" rhs
-      [call]
+      add_stmt <| .set "exception_ty_matches" rhs
     | .none =>
-      [.set "exception_ty_matches" (.boolConst () false)]
+      add_stmt <| .set "exception_ty_matches" (.boolConst () false)
     let cond := .fvar () "exception_ty_matches" none
-    let body_if_matches := body.val.toList.flatMap (PyStmtToBoogie jmp_targets func_infos) ++ [.goto jmp_targets[1]!]
-    set_ex_ty_matches ++ [.ite cond {ss := body_if_matches} {ss := []}]
+    let body_if_matches ← liftPyTrans do
+          body.val.forM (PyStmtToBoogie jmp_targets)
+          add_stmt <| .goto jmp_targets[1]!
+    add_stmt <| .ite cond {ss := body_if_matches} {ss := []}
 
 partial def handleFunctionCall (lhs: List Boogie.Expression.Ident)
-                               (fname: String)
+                               (func : Python.expr SourceRange)
                                (args: Ann (Array (Python.expr SourceRange)) SourceRange)
                                (kwords: Ann (Array (Python.keyword SourceRange)) SourceRange)
                                (_jmp_targets: List String)
-                               (func_infos : List PythonFunctionDecl)
-                               (_s : Python.stmt SourceRange) : List Boogie.Statement :=
+                               (_s : Python.stmt SourceRange) : PyTransM Unit := do
   -- Boogie doesn't allow nested function calls, so we need to introduce temporary variables for each nested call
   let nested_args_calls := args.val.filterMap (λ a => if isCall a then some a else none)
   let args_calls_to_tmps := nested_args_calls.map (λ a => (a, s!"call_arg_tmp_{a.toAst.ann.start}"))
@@ -337,15 +387,22 @@ partial def handleFunctionCall (lhs: List Boogie.Expression.Ident)
 
   let substitution_records : List SubstitutionRecord := args_calls_to_tmps.toList.map (λ p => {pyExpr := p.fst, boogieExpr := .fvar () p.snd none}) ++
                                                         kwords_calls_to_tmps.toList.map (λ p => {pyExpr := p.fst, boogieExpr := .fvar () p.snd none})
-  args_calls_to_tmps.toList.flatMap initTmpParam ++
-    kwords_calls_to_tmps.toList.flatMap initTmpParam ++
-    [.call lhs fname (argsAndKWordsToCanonicalList func_infos fname args.val kwords.val substitution_records)]
-
-partial def PyStmtToBoogie (jmp_targets: List String) (func_infos : List PythonFunctionDecl) (s : Python.stmt SourceRange) : List Boogie.Statement :=
+  let (fname, r) ←
+    match argsAndKWordsToCanonicalList (←read) func args.val kwords.val substitution_records with
+    | .ok p => pure p
+    | .error msg => throw msg
+
+  for c in args_calls_to_tmps do
+    initTmpParam c |>.forM add_stmt
+  for c in kwords_calls_to_tmps do
+    initTmpParam c |>.forM add_stmt
+  add_stmt <| .call lhs fname r
+
+partial def PyStmtToBoogie (jmp_targets: List String) (s : Python.stmt SourceRange) : PyTransM Unit := do
   assert! jmp_targets.length > 0
-  let non_throw := match s with
+  match s with
     | .Import _ names =>
-      [.call [] "import" [PyListStrToBoogie names.val]]
+      add_stmt <| .call [] "import" [PyListStrToBoogie names.val]
     | .ImportFrom _ s names i =>
       let n := match s.val with
       | some s => [strToBoogieExpr s.val]
@@ -353,67 +410,74 @@ partial def PyStmtToBoogie (jmp_targets: List String) (func_infos : List PythonF
       let i := match i.val with
       | some i => [intToBoogieExpr (PyIntToInt i)]
       | none => []
-      [.call [] "importFrom" (n ++ [PyListStrToBoogie names.val] ++ i)]
+      add_stmt <| .call [] "importFrom" (n ++ [PyListStrToBoogie names.val] ++ i)
     | .Expr _ (.Call _ func args kwords) =>
-      let fname := PyExprToString func
-      if callCanThrow func_infos s then
-        handleFunctionCall ["maybe_except"] fname args kwords jmp_targets func_infos s
+      if callCanThrow (←read) s then
+        handleFunctionCall ["maybe_except"] func args kwords jmp_targets s
       else
-        handleFunctionCall [] fname args kwords jmp_targets func_infos s
+        handleFunctionCall [] func args kwords jmp_targets s
     | .Expr _ _ =>
       panic! "Can't handle Expr statements that aren't calls"
     | .Assign _ lhs (.Call _ func args kwords) _ =>
       assert! lhs.val.size == 1
-      let fname := PyExprToString func
-      handleFunctionCall [PyExprToString lhs.val[0]!, "maybe_except"] fname args kwords jmp_targets func_infos s
+      handleFunctionCall [PyExprToString lhs.val[0]!, "maybe_except"] func args kwords jmp_targets s
     | .Assign _ lhs rhs _ =>
       assert! lhs.val.size == 1
-      [.set (PyExprToString lhs.val[0]!) (PyExprToBoogie rhs)]
+      add_stmt <| .set (PyExprToString lhs.val[0]!) (PyExprToBoogie rhs)
     | .AnnAssign _ lhs _ { ann := _ , val := (.some (.Call _ func args kwords))} _ =>
-      let fname := PyExprToString func
-      handleFunctionCall [PyExprToString lhs, "maybe_except"] fname args kwords jmp_targets func_infos s
+      handleFunctionCall [PyExprToString lhs, "maybe_except"] func args kwords jmp_targets s
     | .AnnAssign _ lhs _ { ann := _ , val := (.some (.ListComp _ _ _))} _ =>
       -- TODO: check for errors first
-      [.havoc (PyExprToString lhs)]
+      add_stmt <| .havoc (PyExprToString lhs)
     | .AnnAssign _ lhs _ {ann := _, val := (.some e)} _ =>
-      [.set (PyExprToString lhs) (PyExprToBoogie e)]
+      add_stmt <| .set (PyExprToString lhs) (PyExprToBoogie e)
     | .Try _ body handlers _orelse _finalbody =>
         let new_target := s!"excepthandlers_{jmp_targets[0]!}"
-        let entry_except_handlers := [.block new_target {ss := []}]
         let new_jmp_stack := new_target :: jmp_targets
-        let except_handlers := handlers.val.toList.flatMap (exceptHandlersToBoogie new_jmp_stack func_infos)
-        let var_decls := collectVarDecls body.val
-        [.block "try_block" {ss := var_decls ++ body.val.toList.flatMap (PyStmtToBoogie new_jmp_stack func_infos) ++ entry_except_handlers ++ except_handlers}]
+        let body_stmts ← liftPyTrans do
+              collectVarDecls body.val |>.forM add_stmt
+              body.val.forM (PyStmtToBoogie new_jmp_stack)
+              add_stmt <| .block new_target {ss := []}
+              handlers.val.forM (exceptHandlersToBoogie new_jmp_stack)
+        add_stmt <| .block "try_block" {ss := body_stmts}
     | .FunctionDef _ _ _ _ _ _ _ _ => panic! "Can't translate FunctionDef to Boogie statement"
     | .If _ test then_b else_b =>
-      [.ite (PyExprToBoogie test) {ss := (ArrPyStmtToBoogie func_infos then_b.val)} {ss := (ArrPyStmtToBoogie func_infos else_b.val)}] -- TODO: fix this
+      add_stmt <| .ite (PyExprToBoogie test)
+        (← ArrPyStmtToBoogieM then_b.val)
+        (← ArrPyStmtToBoogieM else_b.val) -- TODO: fix this
     | .Return _ v =>
       match v.val with
-      | .some v => [.set "ret" (PyExprToBoogie v), .goto jmp_targets[0]!] -- TODO: need to thread return value name here. For now, assume "ret"
-      | .none => [.goto jmp_targets[0]!]
+      | .some v =>
+        add_stmt <| .set "ret" (PyExprToBoogie v)
+        add_stmt <| .goto jmp_targets[0]! -- TODO: need to thread return value name here. For now, assume "ret"
+      | .none =>
+        add_stmt <| .goto jmp_targets[0]!
     | .For _ _tgt itr body _ _ =>
       -- Do one unrolling:
       let guard := .app () (.op () "Bool.Not" none) (.eq () (.app () (.op () "dict_str_any_length" none) (PyExprToBoogie itr)) (.intConst () 0))
       -- let guard := .boolConst () true
-      [.ite guard {ss := (ArrPyStmtToBoogie func_infos body.val)} {ss := []}]
+      add_stmt <| .ite guard (← ArrPyStmtToBoogieM body.val) {ss := []}
       -- TODO: missing havoc
     | _ =>
       panic! s!"Unsupported {repr s}"
-  if callCanThrow func_infos s then
-    non_throw ++ [handleCallThrow jmp_targets[0]!]
-  else
-    non_throw
+  if callCanThrow (←read) s then
+    add_stmt <| handleCallThrow jmp_targets[0]!
+
+partial def ArrPyStmtToBoogieM (a : Array (Python.stmt SourceRange)) : PyTransM Boogie.Block := do
+  return { ss := ← liftPyTrans (a.forM (PyStmtToBoogie ["end"])) }
 
-partial def ArrPyStmtToBoogie (func_infos : List PythonFunctionDecl) (a : Array (Python.stmt SourceRange)) : List Boogie.Statement :=
-  a.toList.flatMap (PyStmtToBoogie ["end"] func_infos)
+partial def leanBug : List Nat := []
 
 end --mutual
 
 
 
-def translateFunctions (a : Array (Python.stmt SourceRange)) (func_infos : List PythonFunctionDecl) : List Boogie.Decl :=
-  a.toList.filterMap (λ s => match s with
-    | .FunctionDef _ name _args body _ _ret _ _ =>
+partial def ArrPyStmtToBoogie (ctx : PyTransContext) (a : Array (Python.stmt SourceRange)) : Except ErrorMsg (List Boogie.Statement) :=
+  runPyTrans ctx <| a.forM (PyStmtToBoogie ["end"])
+
+def translateFunctions (a : Array (Python.stmt SourceRange)) (ctx : PyTransContext) : Except ErrorMsg (List Boogie.Decl) :=
+  a.toList.filterMapM λ s => match s with
+    | .FunctionDef _ name _args body _ _ret _ _ => do
 
       let varDecls : List Boogie.Statement := []
       let proc : Boogie.Procedure := {
@@ -423,25 +487,25 @@ def translateFunctions (a : Array (Python.stmt SourceRange)) (func_infos : List
                inputs := [],
                outputs := [("maybe_except", (.tcons "ExceptOrNone" []))]},
         spec := default,
-        body := varDecls ++ ArrPyStmtToBoogie func_infos body.val ++ [.block "end" {ss := []}]
+        body := varDecls ++ (← ArrPyStmtToBoogie ctx body.val) ++ [.block "end" {ss := []}]
       }
-      some (.proc proc)
-    | _ => none)
+      return some (.proc proc)
+    | _ => return none
 
 def pyTyStrToLMonoTy (ty_str: String) : Lambda.LMonoTy :=
   match ty_str with
   | "str" => mty[string]
   | _ => panic! s!"Unsupported type: {ty_str}"
 
-def pythonFuncToBoogie (name : String) (args: List (String × String)) (body: Array (Python.stmt SourceRange)) (ret : Option (Python.expr SourceRange)) (spec : Boogie.Procedure.Spec) (func_infos : List PythonFunctionDecl) : Boogie.Procedure :=
+def pythonFuncToBoogie (name : String) (args: List (String × String)) (body: Array (Python.stmt SourceRange)) (ret : Option (Python.expr SourceRange)) (spec : Boogie.Procedure.Spec) (ctx : PyTransContext) : Except ErrorMsg Boogie.Procedure := do
   let inputs : List (Lambda.Identifier Boogie.Visibility × Lambda.LMonoTy) := args.map (λ p => (p.fst, pyTyStrToLMonoTy p.snd))
   let varDecls := collectVarDecls body ++ [(.init "exception_ty_matches" t[bool] (.boolConst () false)), (.havoc "exception_ty_matches")]
-  let stmts := ArrPyStmtToBoogie func_infos body
+  let stmts ← ArrPyStmtToBoogie ctx body
   let body := varDecls ++ stmts ++ [.block "end" {ss := []}]
   let outputs : Lambda.LMonoTySignature := match ret with
-  | .some v => [("ret", (.tcons "DictStrAny" [])), ("maybe_except", (.tcons "ExceptOrNone" []))]
+  | .some _ => [("ret", (.tcons "DictStrAny" [])), ("maybe_except", (.tcons "ExceptOrNone" []))]
   | .none => [("maybe_except", (.tcons "ExceptOrNone" []))]
-  {
+  return {
     header := {name,
                typeArgs := [],
                inputs,
@@ -462,14 +526,14 @@ def unpackPyArguments (args: Python.arguments SourceRange) : List (String × Str
       | .some ty => (name.val, PyExprToString ty)
       | _ => panic! s!"Missing type annotation on arg: {repr a} ({repr args})")
 
-def PyFuncDefToBoogie (s: Python.stmt SourceRange) (func_infos : List PythonFunctionDecl) : Boogie.Decl × PythonFunctionDecl :=
+def PyFuncDefToBoogie (s: Python.stmt SourceRange) (ctx : PyTransContext) : Except ErrorMsg (Boogie.Decl × PythonFunctionDecl) :=
   match s with
-  | .FunctionDef _ name args body _ ret _ _ =>
+  | .FunctionDef _ name args body _ ret _ _ => do
     let args := unpackPyArguments args
-    (.proc (pythonFuncToBoogie name.val args body.val ret.val default func_infos), {name := name.val, args})
+    pure (.proc (← pythonFuncToBoogie name.val args body.val ret.val default ctx), {name := name.val, args})
   | _ => panic! s!"Expected function def: {repr s}"
 
-def pythonToBoogie (pgm: Strata.Program): Boogie.Program :=
+def pythonToBoogie (sig_db : Python.FuncSigDB) (pgm: Strata.Program): Boogie.Program × Array ErrorMsg :=
   let pyCmds := toPyCommands pgm.commands
   assert! pyCmds.size == 1
   let insideMod := unwrapModule pyCmds[0]!
@@ -481,22 +545,32 @@ def pythonToBoogie (pgm: Strata.Program): Boogie.Program :=
   | .FunctionDef _ _ _ _ _ _ _ _ => false
   | _ => true)
 
-  let globals := [(.var "__name__" (.forAll [] mty[string]) (.strConst () "__main__"))]
+  let globals := #[(.var "__name__" (.forAll [] mty[string]) (.strConst () "__main__"))]
 
-  let rec helper (f : Python.stmt SourceRange → List PythonFunctionDecl → Boogie.Decl × PythonFunctionDecl)
-                 (acc : List PythonFunctionDecl) :
-                 List (Python.stmt SourceRange) → List Boogie.Decl × List PythonFunctionDecl
-  | [] => ([], acc)
+  let rec helper (f : Python.stmt SourceRange → PyTransContext → Except ErrorMsg (Boogie.Decl × PythonFunctionDecl))
+                 (decls : Array Boogie.Decl)
+                 (acc : PyTransContext)
+                 (errors : Array ErrorMsg) :
+                 List (Python.stmt SourceRange) → (Array Boogie.Decl × PyTransContext × Array ErrorMsg)
+  | [] =>
+    (decls, acc, errors)
   | x :: xs =>
-    let (y, acc') := f x acc
-    let new_acc := acc' :: acc
-    let (ys, acc'') := helper f new_acc xs
-    (y :: ys, acc'')
-
-  let func_defs_and_infos := (helper PyFuncDefToBoogie [] func_defs.toList)
-  let func_defs := func_defs_and_infos.fst
-  let func_infos := func_defs_and_infos.snd
+    let ctx := acc
+    match f x ctx with
+    | .ok (y, new_fn) =>
+      let new_ctx :=  ctx.add! new_fn
+      helper f (decls.push y) new_ctx errors xs
+    | .error msg =>
+      helper f decls ctx (errors.push msg) xs
+
+  let (decls, ctx, errors) := helper PyFuncDefToBoogie globals { sig_db := sig_db } #[] func_defs.toList
+  let (decls, errors) :=
+        match pythonFuncToBoogie "__main__" [] non_func_blocks none default ctx with
+        | .ok d => (decls.push (.proc d), errors)
+        | .error msg => (decls, errors.push msg)
+  ({decls := decls.toList}, errors)
+
+end Python
 
-  {decls := globals ++ func_defs ++ [.proc (pythonFuncToBoogie "__main__" [] non_func_blocks none default func_infos)]}
 
 end Strata
diff --git a/StrataMain.lean b/StrataMain.lean
index 28b30d303..ba1cbbf2a 100644
--- a/StrataMain.lean
+++ b/StrataMain.lean
@@ -159,37 +159,59 @@ def diffCommand : Command where
     | _, _ =>
       exitFailure "Cannot compare dialect def with another dialect/program."
 
-def readPythonStrata (path : String) : IO Strata.Program := do
-  let bytes ← Strata.Util.readBinInputSource path
+def readPythonStrata (strata_path : String) : IO Strata.Program := do
+  let bytes ← Strata.Util.readBinInputSource strata_path
   if ! bytes.startsWith Ion.binaryVersionMarker then
     exitFailure s!"pyAnalyze expected Ion file"
   match Strata.Program.fromIon Strata.Python.Python_map Strata.Python.Python.name bytes with
   | .ok p => pure p
   | .error msg => exitFailure msg
 
+def pyTranslate (py_path : String) (fm : Lean.FileMap) (pgm : Strata.Program) : IO Boogie.Program := do
+  let preludePgm := Strata.Python.Internal.Boogie.prelude
+  let (bpgm, errors) := Strata.Python.pythonToBoogie Strata.Python.funcSigs pgm
+  if errors.size > 0 then
+    IO.eprintln s!"Errors in translation"
+    for e in errors do
+      let spos := fm.toPosition e.loc.start
+      let epos := fm.toPosition e.loc.stop
+      -- Render the error location information in a format VSCode understands
+      let pos :=
+            if spos.line == spos.line then
+              s!"{py_path}:{spos.line}:{spos.column+1}-{epos.column+1}"
+            else
+              s!"{py_path}:{spos.line}:{spos.column+1}"
+
+      IO.eprintln s!"{pos}: {e.message}"
+    IO.Process.exit 1
+  pure { decls := preludePgm.decls ++ bpgm.decls }
+
 def pyTranslateCommand : Command where
   name := "pyTranslate"
-  args := [ "file" ]
+  args := [ "py_file", "strata_file" ]
   help := "Translate a Strata Python Ion file to Strata.Boogie. Write results to stdout."
   callback := fun _ v => do
-    let pgm ← readPythonStrata v[0]
-    let preludePgm := Strata.Python.Internal.Boogie.prelude
-    let bpgm := Strata.pythonToBoogie pgm
-    let newPgm : Boogie.Program := { decls := preludePgm.decls ++ bpgm.decls }
+    let python_file := v[0]
+    let strata_file := v[1]
+    let fm : Lean.FileMap := .ofString (← IO.FS.readFile python_file)
+    let pgm ← readPythonStrata strata_file
+    let newPgm : Boogie.Program ← pyTranslate python_file fm pgm
     IO.print newPgm
 
 def pyAnalyzeCommand : Command where
   name := "pyAnalyze"
-  args := [ "file", "verbose" ]
+  args := [ "py_file", "strata_file", "verbose" ]
   help := "Analyze a Strata Python Ion file. Write results to stdout."
   callback := fun _ v => do
-    let verbose := v[1] == "1"
-    let pgm ← readPythonStrata v[0]
+    let python_file := v[0]
+    let strata_file := v[1]
+    let verbose := v[2] == "1"
+    let fm : Lean.FileMap := .ofString (← IO.FS.readFile python_file)
+
+    let pgm ← readPythonStrata strata_file
     if verbose then
       IO.print pgm
-    let preludePgm := Strata.Python.Internal.Boogie.prelude
-    let bpgm := Strata.pythonToBoogie pgm
-    let newPgm : Boogie.Program := { decls := preludePgm.decls ++ bpgm.decls }
+    let newPgm : Boogie.Program ← pyTranslate python_file fm pgm
     if verbose then
       IO.print newPgm
     let vcResults ← EIO.toIO (fun f => IO.Error.userError (toString f))