PythonToCore.lean

/-
  Copyright Strata Contributors

  SPDX-License-Identifier: Apache-2.0 OR MIT
-/
module
public import StrataPython.PythonDialect
public import StrataPython.FunctionSignatures
public import Strata.Languages.Core.Program
import Strata.Languages.Core.Env

open Strata
open StrataDDM

namespace StrataPython
open Lambda.LTy.Syntax

public section

-- Some hard-coded things we'll need to fix later:

def clientType : Core.Expression.Ty := .forAll [] (.tcons "Client" [])
def dummyClient : Core.Expression.Expr := .fvar () "DUMMY_CLIENT" none

def dictStrAnyType : Core.Expression.Ty := .forAll [] (.tcons "DictStrAny" [])
def dummyDictStrAny : Core.Expression.Expr := .fvar () "DUMMY_DICT_STR_ANY" none

def strType : Core.Expression.Ty := .forAll [] (.tcons "string" [])
def dummyStr : Core.Expression.Expr := .fvar () "DUMMY_STR" none

def listStrType : Core.Expression.Ty := .forAll [] (.tcons "ListStr" [])
def dummyListStr : Core.Expression.Expr := .fvar () "DUMMY_LIST_STR" none

def datetimeType : Core.Expression.Ty := .forAll [] (.tcons "Datetime" [])
def dummyDatetime : Core.Expression.Expr := .fvar () "DUMMY_DATETIME" none

def dateType : Core.Expression.Ty := .forAll [] (.tcons "Date" [])
def dummyDate : Core.Expression.Expr := .fvar () "DUMMY_DATE" none

def timedeltaType : Core.Expression.Ty := .forAll [] (.tcons "int" [])
def dummyTimedelta : Core.Expression.Expr := .fvar () "DUMMY_Timedelta" none

-------------------------------------------------------------------------------

-- Translating a Python expression can require Strata Core statements, e.g., a function call
-- We translate these by first defining temporary variables to store the results of the stmts
-- and then using those variables in the expression.
structure PyExprTranslated where
  stmts : List Core.Statement
  expr: Core.Expression.Expr
  post_stmts : List Core.Statement := []
deriving Inhabited

structure PythonFunctionDecl where
  name : String
  args : List (String × String) -- Elements are (arg_name, arg_ty) where `arg_ty` is the string representation of the type in Python
  ret : String
deriving Repr, BEq, Inhabited

structure PythonClassDecl where
  name : String
deriving Repr, BEq, Inhabited

inductive ExtrinsicModelingChoice where
  | havocAll : ExtrinsicModelingChoice
  | havocArgsAndRet : ExtrinsicModelingChoice
  | havocRet : ExtrinsicModelingChoice
deriving Repr, BEq, Inhabited

structure ExtrinsicsModelingConfig where
  behaviors : String → ExtrinsicModelingChoice

def defaultExtrinsicsModelChoice (_s: String) : ExtrinsicModelingChoice :=
  .havocArgsAndRet

instance : Inhabited ExtrinsicsModelingConfig where
  default := {behaviors := defaultExtrinsicsModelChoice}

def extractProcedureNames (pgm : Core.Program) : List String :=
  pgm.decls.filterMap fun decl =>
    match decl with
    | .proc p _ => some p.header.name.name
    | _ => none

structure TranslationContext where
  signatures : Signatures
  filePath : String := ""
  expectedType : Option (Lambda.LMonoTy) := none
  variableTypes : List (String × Lambda.LMonoTy) := []
  func_infos : List PythonFunctionDecl := []
  class_infos : List PythonClassDecl := []
  extrinsicsModelConfig : ExtrinsicsModelingConfig := default
  preludeProcedureNames : List String := []
deriving Inhabited

/-- Create metadata from a SourceRange for attaching to Core statements. -/
def sourceRangeToMetaData (filePath : String) (sr : SourceRange) : Imperative.MetaData Core.Expression :=
  Imperative.MetaData.ofSourceRange (.file filePath) sr

-------------------------------------------------------------------------------

def strToCoreExpr (s: String) : Core.Expression.Expr :=
  .strConst () s

def intToCoreExpr (i: Int) : Core.Expression.Expr :=
  .intConst () i

def PyIntToInt (i : StrataPython.int SourceRange) : Int :=
  match i with
  | .IntPos _ n => n.val
  | .IntNeg _ n => -n.val

def PyConstToCore (c: constant SourceRange) : Core.Expression.Expr :=
  match c with
  | .ConString _ s => .strConst () s.val
  | .ConPos _ i => .intConst () i.val
  | .ConNeg _ i => .intConst () (-i.val)
  | .ConBytes _ _b => .const () (.strConst "") -- TODO: fix
  | .ConFloat _ f => .strConst () (f.val)
  | _ => panic! s!"Unhandled Constant: {repr c}"

def PyAliasToCoreExpr (a : alias SourceRange) : Core.Expression.Expr :=
  match a with
  | .mk_alias _ n as_n =>
  assert! as_n.val.isNone
  .strConst () n.val

def handleAdd (translation_ctx: TranslationContext) (lhs rhs: Core.Expression.Expr) : Core.Expression.Expr :=
  match lhs, rhs with
  | .intConst () l, .intConst () r => .intConst () (l + r)
  | .fvar () l _, .fvar () r _ =>
    let l_ty := translation_ctx.variableTypes.find? (λ p => p.fst == l.name)
    let r_ty := translation_ctx.variableTypes.find? (λ p => p.fst == r.name)
    match l_ty, r_ty with
    | some (_, .tcons "int" []), some (_, .tcons "int" []) =>
      .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Add⟩) (some mty[int → (int → int)])) lhs) rhs
    | some (_, .tcons "string" []), some (_, .tcons "string" []) =>
      .app () (.app () (.op () "Str.Concat" mty[string → (string → string)]) lhs) rhs
    | _, _ => panic! s!"Unsupported types for +. Exprs: {lhs} and {rhs}"
  | _, _ => .app () (.app () (.op () "Str.Concat" mty[string → (string → string)]) lhs) rhs

def handleSub (translation_ctx: TranslationContext) (lhs rhs: Core.Expression.Expr) : Core.Expression.Expr :=
  match lhs, rhs with
  | .intConst () l, .intConst () r => .intConst () (l - r)
  | .fvar () l _, .fvar () r _ =>
    let l_ty := translation_ctx.variableTypes.find? (λ p => p.fst == l.name)
    let r_ty := translation_ctx.variableTypes.find? (λ p => p.fst == r.name)
    match l_ty, r_ty with
    | some (_, .tcons "int" []), some (_, .tcons "int" []) =>
      .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Sub⟩) (some mty[int → (int → int)])) lhs) rhs
    | some (_, .tcons "Datetime" []), some (_, .tcons "int" []) =>
      .app () (.app () (.op () "Datetime_sub" none) lhs) rhs
    | some (_, .tcons "Datetime" []), some (_, .tcons "Timedelta" []) =>
      .app () (.app () (.op () "Datetime_sub" none) lhs) rhs
    | _, _ => panic! s!"Unsupported types for -. Exprs: {lhs} and {rhs}"
  | _, _ => panic! s!"Unsupported args for -. Got: {lhs} and {rhs}"

def handleMult (translation_ctx: TranslationContext) (lhs rhs: Core.Expression.Expr) : Core.Expression.Expr :=
  match lhs, rhs with
  | .strConst () s, .intConst () i => .strConst () (String.join (List.replicate i.toNat s))
  | .intConst () l, .intConst () r => .intConst () (l * r)
  | .fvar () l _, .fvar () r _ =>
    let l := translation_ctx.variableTypes.find? (λ p => p.fst == l.name)
    let r := translation_ctx.variableTypes.find? (λ p => p.fst == r.name)
    match l, r with
    | .some lty, .some rty =>
      match lty.snd, rty.snd with
      | .tcons "int" [], .tcons "int" [] => .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Mul⟩) (some mty[int → (int → int)])) lhs) rhs
      | _, _ => panic! s!"Unsupported types for fvar *. Types: {lty} and {rty}"
    | _, _ => panic! s!"Missing needed type information for *. Exprs: {lhs} and {rhs}"
  | _ , _ => panic! s!"Unsupported args for * . Got: {lhs} and {rhs}"

def handleFloorDiv (_translation_ctx: TranslationContext) (lhs rhs: Core.Expression.Expr) : Core.Expression.Expr :=
  .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Div⟩) (some mty[int → (int → int)])) lhs) rhs

def handleNot (arg: Core.Expression.Expr) : Core.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 handleLt (translation_ctx: TranslationContext) (lhs rhs: Core.Expression.Expr) : Core.Expression.Expr :=
  match lhs, rhs with
  | .fvar () l _, .fvar () r _ =>
    let l_ty := translation_ctx.variableTypes.find? (λ p => p.fst == l.name)
    let r_ty := translation_ctx.variableTypes.find? (λ p => p.fst == r.name)
    match l_ty, r_ty with
    | some (_, .tcons "Datetime" []), some (_, .tcons "Datetime" []) =>
      .app () (.app () (.op () "Datetime_lt" none) lhs) rhs
    | _, _ => .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Lt⟩) (some mty[int → (int → bool)])) lhs) rhs
  | _, _ => .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Lt⟩) (some mty[int → (int → bool)])) lhs) rhs

def handleLtE (translation_ctx: TranslationContext) (lhs rhs: Core.Expression.Expr) : Core.Expression.Expr :=
  match lhs, rhs with
  | .fvar () l _, .fvar () r _ =>
    let l_ty := translation_ctx.variableTypes.find? (λ p => p.fst == l.name)
    let r_ty := translation_ctx.variableTypes.find? (λ p => p.fst == r.name)
    match l_ty, r_ty with
    | some (_, .tcons "Datetime" []), some (_, .tcons "Datetime" []) =>
      let eq := (.eq () lhs rhs)
      let lt := (.app () (.app () (.op () "Datetime_lt" none) lhs) rhs)
      (.app () (.app () (Core.coreOpExpr (.bool .Or)) eq) lt)
    | _, _ => .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Le⟩) (some mty[int → (int → bool)])) lhs) rhs
  | _, _ => .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Le⟩) (some mty[int → (int → bool)])) lhs) rhs

def handleGt (lhs rhs: Core.Expression.Expr) : Core.Expression.Expr :=
  .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Gt⟩) (some mty[int → (int → bool)])) lhs) rhs

def handleGtE (lhs rhs: Core.Expression.Expr) : Core.Expression.Expr :=
  .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Ge⟩) (some mty[int → (int → bool)])) lhs) rhs

structure SubstitutionRecord where
  pyExpr : expr SourceRange
  coreExpr : Core.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: expr SourceRange) : Bool :=
  match e1, e2 with
  | .Call sr1 _ _ _, .Call sr2 _ _ _ => sr1 == sr2
  | _ , _ => false

-- TODO: handle rest of names
def PyListStrToCore (names : Array (alias SourceRange)) : Core.Expression.Expr :=
  .app () (.app () (.op () "ListStr_cons" mty[string → (ListStr → ListStr)]) (PyAliasToCoreExpr names[0]!))
       (.op () "ListStr_nil" mty[ListStr])

def handleList (_elmts: Array (expr SourceRange)) (expected_type : Lambda.LMonoTy): PyExprTranslated :=
  match expected_type with
  | (.tcons "ListStr" _) => {stmts := [], expr := (.op () "ListStr_nil" expected_type)}
  | (.tcons "ListDictStrAny" _) => {stmts := [], expr := (.op () "ListDictStrAny_nil" expected_type)}
  | _ => panic! s!"Unexpected type : {expected_type}"

def PyOptExprToString (e : opt_expr SourceRange) : String :=
  match e with
  | .some_expr _ (.Constant _ (.ConString _ s) _) => s.val
  | _ => panic! s!"Expected some constant string: {repr e}"

partial def PyExprToString (e : expr SourceRange) : String :=
  match e with
  | .Name _ n _ => n.val
  | .Attribute _ v attr _ => s!"{PyExprToString v}_{attr.val}"
  | .Subscript _ v slice _ =>
    let v_name := PyExprToString v
    match v_name with
    | "Dict" =>
      match slice with
      | .Tuple _ elts _ =>
        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}"

def PyExprToMonoTy (e : expr SourceRange) : Lambda.LMonoTy :=
  match e with
  | .Name _ n _ =>
    match n.val with
    | "str" => .tcons "string" []
    | "float" => .tcons "string" []
    | "Dict[str Any]" => .tcons "DictStrAny" []
    | "List[str]" => .tcons "ListStr" []
    | "datetime" => .tcons "Datetime" []
    | "date" => .tcons "Date" []
    | "timedelta" => .tcons "Timedelta" []
    | _ => .tcons n.val []
  | .Subscript _ val _slice _ =>
    match val with
    | .Name _ n _ =>
      match n.val with
      | "Dict" => .tcons "DictStrAny" []
      | "List" => .tcons "ListStr" []
      | _ => panic! s!"Unsupported name: {repr n}"
    | _ => panic! s!"Expected name: {repr e}"
  | _ => panic! s!"Unhandled Expr: {repr e}"


-- 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: 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)
    | _ => false
  | _ => false

def noneOrExpr (translation_ctx : TranslationContext) (fname n : String) (e: Core.Expression.Expr) : Core.Expression.Expr :=
  let type_str := translation_ctx.signatures.getFuncSigType fname n
  match type_str with
  | .error e => panic! s!"noneOrExpr called without type_str: {e}"
  | .ok type_str =>
    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

def handleCallThrow (jmp_target : String) : Core.Statement :=
  let cond := .app () (.op () "ExceptOrNone..isExceptOrNone_mk_code" none) (.fvar () "maybe_except" none)
  .ite (.det cond) [.exit jmp_target .empty] [] .empty

def deduplicateTypeAnnotations (l : List (String × Option String)) : List (String × String) := Id.run do
  let mut m : Map String String := []
  for p in l do
    let name := p.fst
    let oty := p.snd
    match oty with
    | .some ty =>
      match m.find? name with
      | .some other_ty =>
        if ty != other_ty then
          panic! s!"Type annotation mismatch: {other_ty} vs {ty}"
      | .none => m := (name, ty) :: m
    | .none => ()
  let names := l.map (λ p => p.fst)
  let unique_names := names.dedup
  unique_names.map (λ n =>
    match m.find? n with
    | .some ty => (n, ty)
    | .none => panic s!"Missing type annotations for {n}")

partial def collectVarDecls (translation_ctx : TranslationContext) (stmts: Array (stmt SourceRange)) : List Core.Statement :=
  let rec go (s : 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
    | .For _ _ _ body _ _ => body.val.toList.flatMap go
    | _ => []
  let dup := stmts.toList.flatMap go
  let dedup := deduplicateTypeAnnotations dup
  let toCore (p: String × String) : List Core.Statement :=
    let name := p.fst
    let ty_name := p.snd
    match ty_name with
    | "bool" => [(.init name t[bool] (.det (.boolConst () false)) .empty), (.havoc name .empty)]
    | "str" => [(.init name t[string] (.det (.strConst () "")) .empty), (.havoc name .empty)]
    | "int" => [(.init name t[int] (.det (.intConst () 0)) .empty), (.havoc name .empty)]
    | "float" => [(.init name t[string] (.det (.strConst () "0.0")) .empty), (.havoc name .empty)] -- Floats as strs for now
    | "bytes" => [(.init name t[string] (.det (.strConst () "")) .empty), (.havoc name .empty)]
    | "Client" => [(.init name clientType (.det dummyClient) .empty), (.havoc name .empty)]
    | "Dict[str Any]" => [(.init name dictStrAnyType (.det dummyDictStrAny) .empty), (.havoc name .empty)]
    | "List[str]" => [(.init name listStrType (.det dummyListStr) .empty), (.havoc name .empty)]
    | "datetime" => [(.init name datetimeType (.det dummyDatetime) .empty), (.havoc name .empty)]
    | "date" => [(.init name dateType (.det dummyDate) .empty), (.havoc name .empty)]
    | "timedelta" => [(.init name timedeltaType (.det dummyTimedelta) .empty), (.havoc name .empty)]
    | _ =>
      let user_defined_class := translation_ctx.class_infos.find? (λ i => i.name == ty_name)
      match user_defined_class with
      | .some i =>
        let user_defined_class_ty := .forAll [] (.tcons i.name [])
        let user_defined_class_dummy := .fvar () ("DUMMY_" ++ i.name) none
        [(.init name user_defined_class_ty (.det user_defined_class_dummy) .empty), (.havoc name .empty)]
      | .none => panic! s!"Unsupported type annotation: `{ty_name}`"
  let foo := dedup.map toCore
  foo.flatten

def isCall (e: expr SourceRange) : Bool :=
  match e with
  | .Call _ _ _ _ => true
  | _ => false

def remapFname (translation_ctx: TranslationContext) (fname: String) : String :=
  match translation_ctx.class_infos.find? (λ i => i.name == fname) with
  | .some i =>
    i.name ++ "___init__"
  | _ =>
    match fname with
    | "float" => "str_to_float"
    | _ => fname

def noFuncModel (translation_ctx: TranslationContext) (fname: String) : Bool :=
  let funcInPrelude := translation_ctx.preludeProcedureNames.contains fname
  let funcInSigList := (translation_ctx.signatures.getFuncSigOrder fname).isOk
  let funcInUserCode := translation_ctx.func_infos.any (λ e => e.name == fname) || translation_ctx.class_infos.any (λ e => e.name++"___init__" == fname)
  !(funcInPrelude || funcInSigList || funcInUserCode)

def handleUnmodeledFunCall (lhs: List Core.Expression.Ident)
                               (fname: String)
                               (args: Ann (Array (expr SourceRange)) SourceRange)
                               (kwords: Ann (Array (keyword SourceRange)) SourceRange)
                               (_jmp_targets: List String)
                               (translation_ctx: TranslationContext)
                               (s : stmt SourceRange) : List Core.Statement :=
  let md := sourceRangeToMetaData translation_ctx.filePath s.toAst.ann
  match translation_ctx.extrinsicsModelConfig.behaviors fname with
  | .havocAll => panic! "Unimplemented"
  | .havocArgsAndRet =>
    let nameArgs := args.val.filterMap (λ e => match e with
    | .Name _ n _ => .some n
    | _ => .none)
    let nameKWs := kwords.val.filterMap (λ e => match e with
    | .mk_keyword _ _ expr => match expr with
      | .Name _ n _ => .some n
      | _ => .none)
    let havocArgs := nameArgs.toList.map (λ n => .havoc n.val md) ++ nameKWs.toList.map (λ n => .havoc n.val md)
    lhs.map (λ e => .havoc e md) ++ havocArgs
  | .havocRet =>
    lhs.map (λ e => .havoc e md)

mutual

partial def PyExprToCoreWithSubst (translation_ctx : TranslationContext)  (substitution_records : Option (List SubstitutionRecord)) (e : expr SourceRange) : PyExprTranslated :=
  PyExprToCore translation_ctx e substitution_records

partial def PyKWordsToCore (substitution_records : Option (List SubstitutionRecord)) (kw : keyword SourceRange) : (String × PyExprTranslated) :=
  match kw with
  | .mk_keyword _ name expr =>
    match name.val with
    | some n => (n.val, PyExprToCoreWithSubst default substitution_records expr)
    | none => panic! "Keyword arg should have a name"

-- TODO: we should be checking that args are right
partial def argsAndKWordsToCanonicalList (translation_ctx : TranslationContext)
                                 (fname: String)
                                 (args : Array (expr SourceRange))
                                 (kwords: Array (keyword SourceRange))
                                 (substitution_records : Option (List SubstitutionRecord) := none) : List Core.Expression.Expr × List Core.Statement :=
  if translation_ctx.func_infos.any (λ e => e.name == fname) || translation_ctx.class_infos.any (λ e => e.name++"___init__" == fname) then
    if translation_ctx.func_infos.any (λ e => e.name == fname) then
      (args.toList.map (λ a => (PyExprToCoreWithSubst default substitution_records a).expr), [])
    else
      (args.toList.map (λ a => (PyExprToCoreWithSubst default substitution_records a).expr), [])
  else
    let required_order := translation_ctx.signatures.getFuncSigOrder fname
    match required_order with
    | .error e => panic! s!"argsAndKWordsToCanonicalList called fname not found in getFuncSigOrder: {e}"
    | .ok required_order =>
    assert! args.size <= required_order.length
    let remaining := required_order.drop args.size
    let kws_and_exprs := kwords.toList.map (PyKWordsToCore substitution_records)
    let ordered_remaining_args := remaining.map (λ n => match kws_and_exprs.find? (λ p => p.fst == n) with
      | .some p =>
        noneOrExpr translation_ctx fname n p.snd.expr
      | .none =>
        match translation_ctx.signatures.getFuncSigType fname n with
        | .error e => panic! s!"argsAndKWordsToCanonicalList called fname not found in getFuncSigType: {e}"
        | .ok arg_ty => TypeStrToCoreExpr arg_ty)
    let args := args.map (PyExprToCoreWithSubst default substitution_records)
    let args := (List.range required_order.length).filterMap (λ n =>
        if n < args.size then
          let arg_name := required_order[n]! -- Guaranteed by range. Using finRange causes breaking coercions to Nat.
          some (noneOrExpr translation_ctx fname arg_name args[n]!.expr)
        else
          none)
    (args ++ ordered_remaining_args, kws_and_exprs.flatMap (λ p => p.snd.stmts))

partial def handleDict (translation_ctx: TranslationContext) (sr : SourceRange) (keys: Array (opt_expr SourceRange)) (values: Array (expr SourceRange)) : PyExprTranslated :=
  let md := sourceRangeToMetaData translation_ctx.filePath sr
  let dict := .app () (.op () "DictStrAny_mk" none) (.strConst () "DefaultDict") -- TODO: need to generate unique dict arg
  assert! keys.size == values.size
  let zipped := Array.zip keys values

  let res := zipped.toList.flatMap (λ (k, v) =>
    let n := PyOptExprToString k
    let in_dict := (.assume s!"assume_{n}_in_dict" (.app () (.app () (.op () "str_in_dict_str_any" none) (.strConst () n)) dict) md)
    match v with
    | .Call _ f args _ =>
      match f with
      | .Name _ {ann := _ , val := "str"} _ =>
        assert! args.val.size == 1
        let dt := (.app () (.op () "datetime_to_str" none) ((PyExprToCore default args.val[0]!).expr))
        let dict_of_v_is_k := (.assume s!"assume_{n}_key" (.eq () (.app () (.app () (.op () "dict_str_any_get_str" none) dict) (.strConst () n)) dt) md)
        [in_dict, dict_of_v_is_k]
      | _ => panic! "Unsupported function when constructing map"
    | _ =>
      let dict_of_v_is_k := (.assume s!"assume_{n}_key" (.eq () (.app () (.app () (.op () "dict_str_any_get_str" none) dict) (.strConst () n)) (.strConst () "DummyVal")) md)
      [in_dict, dict_of_v_is_k])

  {stmts := res , expr := dict, post_stmts := []}

partial def PyExprToCore (translation_ctx : TranslationContext) (e : expr SourceRange) (substitution_records : Option (List SubstitutionRecord) := none) : PyExprTranslated :=
  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
    {stmts := [], expr := record.coreExpr}
  else
    match e with
    | .Call _ f args kwords =>
      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 _ => {stmts := [], expr :=  PyConstToCore c}
    | .Name _ n _ =>
      match n.val with
      | "AssertionError" | "Exception" => {stmts := [], expr := .strConst () n.val}
      | s =>
        match translation_ctx.variableTypes.find? (λ p => p.fst == s) with
        | .some p =>
          if translation_ctx.expectedType == some (.tcons "bool" []) && p.snd == (.tcons "DictStrAny" []) then
            let a := .fvar () n.val none
            let e := .app () (Core.coreOpExpr (.bool .Not)) (.eq () (.app () (.op () "dict_str_any_length" none) a) (.intConst () 0))
            {stmts := [], expr := e}
          else
            {stmts := [], expr := .fvar () n.val none}
        | .none => {stmts := [], expr := .fvar () n.val none}
    | .JoinedStr _ ss => PyExprToCore translation_ctx ss.val[0]! -- TODO: need to actually join strings
    | .BinOp _ lhs op rhs =>
      let lhs := (PyExprToCore translation_ctx lhs)
      let rhs := (PyExprToCore translation_ctx rhs)
      match op with
      | .Add _ =>
        {stmts := lhs.stmts ++ rhs.stmts, expr := handleAdd translation_ctx lhs.expr rhs.expr}
      | .Sub _ =>
        {stmts := lhs.stmts ++ rhs.stmts, expr := handleSub translation_ctx lhs.expr rhs.expr}
      | .Mult _ =>
        {stmts := lhs.stmts ++ rhs.stmts, expr := handleMult translation_ctx lhs.expr rhs.expr}
      | .FloorDiv _ =>
        {stmts := lhs.stmts ++ rhs.stmts, expr := handleFloorDiv translation_ctx lhs.expr rhs.expr}
      | _ => panic! s!"Unhandled BinOp: {repr e}"
    | .Compare _ lhs op rhs =>
      let lhs := PyExprToCore translation_ctx lhs
      assert! rhs.val.size == 1
      let rhs := PyExprToCore translation_ctx rhs.val[0]!
      match op.val with
      | #[v] => match v with
        | cmpop.Eq _ =>
          {stmts := lhs.stmts ++ rhs.stmts, expr := (.eq () lhs.expr rhs.expr)}
        | cmpop.In _ =>
          {stmts := lhs.stmts ++ rhs.stmts, expr := .app () (.app () (.op () "str_in_dict_str_any" none) lhs.expr) rhs.expr}
        | cmpop.Lt _ =>
          {stmts := lhs.stmts ++ rhs.stmts, expr := handleLt translation_ctx lhs.expr rhs.expr}
        | cmpop.LtE _ =>
          {stmts := lhs.stmts ++ rhs.stmts, expr := handleLtE translation_ctx lhs.expr rhs.expr}
        | cmpop.Gt _ =>
          {stmts := lhs.stmts ++ rhs.stmts, expr := handleGt lhs.expr rhs.expr}
        | cmpop.GtE _ =>
          {stmts := lhs.stmts ++ rhs.stmts, expr := handleGtE lhs.expr rhs.expr}
        | _ => panic! s!"Unhandled comparison op: {repr op.val}"
      | _ => panic! s!"Unhandled comparison op: {repr op.val}"
    | .Dict sr keys values =>
      let res := handleDict translation_ctx sr keys.val values.val
      res
    | .ListComp _ keys values => panic! "ListComp must be handled at stmt level"
    | .UnaryOp _ op arg => match op with
      | .Not _ => {stmts := [], expr := handleNot (PyExprToCore translation_ctx arg).expr}
      | _ => panic! s!"Unsupported UnaryOp: {repr e}"
    | .Subscript sr_sub v slice _ =>
      let sub_md := sourceRangeToMetaData translation_ctx.filePath sr_sub
      let l := PyExprToCore translation_ctx v
      let k := PyExprToCore translation_ctx slice
      -- TODO: we need to plumb the type of `v` here
      match s!"{repr l.expr}" with
      | "LExpr.fvar () { name := \"keys\", metadata := () } none" =>
          -- let access_check : Core.Statement := .assert "subscript_bounds_check" (.app () (.app () (.op () "str_in_dict_str_any" none) k.expr) l.expr)
          {stmts := l.stmts ++ k.stmts, expr := .app () (.app () (.op () "list_str_get" none) l.expr) k.expr}
      | "LExpr.fvar () { name := \"blended_cost\", metadata := () } none" =>
          -- let access_check : Core.Statement := .assert "subscript_bounds_check" (.app () (.app () (.op () "str_in_dict_str_any" none) k.expr) l.expr)
          {stmts := l.stmts ++ k.stmts, expr := .app () (.app () (.op () "dict_str_any_get_str" none) l.expr) k.expr}
      | _ =>
        match translation_ctx.expectedType with
        | .some (.tcons "ListStr" []) =>
          let access_check : Core.Statement := .assert "subscript_bounds_check" (.app () (.app () (.op () "str_in_dict_str_any" none) k.expr) l.expr) sub_md
          {stmts := l.stmts ++ k.stmts ++ [access_check], expr := .app () (.app () (.op () "dict_str_any_get_list_str" none) l.expr) k.expr}
        | _ =>
          let access_check : Core.Statement := .assert "subscript_bounds_check" (.app () (.app () (.op () "str_in_dict_str_any" none) k.expr) l.expr) sub_md
          {stmts := l.stmts ++ k.stmts ++ [access_check], expr := .app () (.app () (.op () "dict_str_any_get" none) l.expr) k.expr}
    | .List _ elmts _ =>
      match elmts.val[0]! with
      | .Constant _ expr _ => match expr with
        | .ConString _ s => handleList elmts.val (.tcons "ListStr" [])
        | _ => panic! s!"Expr: {repr expr}"
      | .Dict _ _ _ => handleList elmts.val (.tcons "ListDictStrAny" [])
      | _ => panic! s!"Unexpected element: {repr elmts.val[0]!}"
    | .Attribute _ v attr ctx =>
      let pre := PyExprToString v
      {stmts := [], expr := strToCoreExpr (pre ++ "_" ++ attr.val)}
    | _ => panic! s!"Unhandled Expr: {repr e}"

partial def initTmpParam (translation_ctx: TranslationContext) (p: expr SourceRange × String) : List Core.Statement :=
  let md := sourceRangeToMetaData translation_ctx.filePath p.fst.toAst.ann
  match p.fst with
  | .Call _ f args _ =>
    match f with
    | .Name _ n _ =>
      match n.val with
      | "json_dumps" => [(.init p.snd t[string] (.det (.strConst () "")) md), .call "json_dumps" ([.inArg (.app () (.op () "DictStrAny_mk" none) (.strConst () "DefaultDict")), .inArg (TypeStrToCoreExpr "IntOrNone")] ++ [.outArg p.snd, .outArg "maybe_except"]) md]
      | "str" =>
        assert! args.val.size == 1
        [(.init p.snd t[string] (.det (.strConst () "")) md), .set p.snd (.app () (.op () "datetime_to_str" none) ((PyExprToCore default args.val[0]!).expr)) md]
      | "int" => [(.init p.snd t[int] (.det (.intConst () 0)) md), .set p.snd (.op () "datetime_to_int" none) md]
      | _ => panic! s!"Unsupported name {n.val}"
    | _ => panic! s!"Unsupported tmp param init call: {repr f}"
  | _ => panic! "Expected Call"

partial def exceptHandlersToCore (jmp_targets: List String) (translation_ctx: TranslationContext) (h : excepthandler SourceRange) : List Core.Statement :=
  assert! jmp_targets.length >= 2
  match h with
  | .ExceptHandler sr ex_ty _ body =>
    let md := sourceRangeToMetaData translation_ctx.filePath sr
    let set_ex_ty_matches := match ex_ty.val with
    | .some ex_ty =>
      let inherits_from : Core.CoreIdent := "inheritsFrom"
      let get_ex_tag : Core.CoreIdent := "ExceptOrNone..code_val!"
      let exception_ty : Core.Expression.Expr := .app () (.op () get_ex_tag none) (.fvar () "maybe_except" none)
      let rhs_curried : Core.Expression.Expr := .app () (.op () inherits_from none) exception_ty
      let res := PyExprToCore translation_ctx ex_ty
      let rhs : Core.Expression.Expr := .app () rhs_curried (res.expr)
      let call := .set "exception_ty_matches" rhs md
      res.stmts ++ [call]
    | .none =>
      [.set "exception_ty_matches" (.boolConst () false) md]
    let cond := .fvar () "exception_ty_matches" none
    let body_if_matches := body.val.toList.flatMap (λ s => (PyStmtToCore jmp_targets.tail! translation_ctx s).fst) ++ [.exit jmp_targets[1]! md]
    set_ex_ty_matches ++ [.ite (.det cond) body_if_matches [] md]

partial def handleFunctionCall (lhs: List Core.Expression.Ident)
                               (fname: String)
                               (args: Ann (Array (expr SourceRange)) SourceRange)
                               (kwords: Ann (Array (keyword SourceRange)) SourceRange)
                               (jmp_targets: List String)
                               (translation_ctx: TranslationContext)
                               (s : stmt SourceRange) : List Core.Statement :=

  let fname := remapFname translation_ctx fname

  if noFuncModel translation_ctx fname then
    handleUnmodeledFunCall lhs fname args kwords jmp_targets translation_ctx s
  else

  -- Strata Core 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, coreExpr := .fvar () p.snd none}) ++
                                                        kwords_calls_to_tmps.toList.map (λ p => {pyExpr := p.fst, coreExpr := .fvar () p.snd none})

  let md := sourceRangeToMetaData translation_ctx.filePath s.toAst.ann
  let res := argsAndKWordsToCanonicalList translation_ctx fname args.val kwords.val substitution_records
  args_calls_to_tmps.toList.flatMap (initTmpParam translation_ctx) ++
    kwords_calls_to_tmps.toList.flatMap (initTmpParam translation_ctx) ++
    res.snd ++ [.call fname (res.fst.map .inArg ++ lhs.map .outArg) md]

partial def handleComprehension (translation_ctx: TranslationContext) (lhs: expr SourceRange) (gen: Array (comprehension SourceRange)) : List Core.Statement :=
  assert! gen.size == 1
  match gen[0]! with
  | .mk_comprehension sr _ itr _ _ =>
    let md := sourceRangeToMetaData translation_ctx.filePath sr
    let res := PyExprToCore default itr
    let guard := .app () (Core.coreOpExpr (.bool .Not)) (.eq () (.app () (.op () "dict_str_any_length" none) res.expr) (.intConst () 0))
    let then_ss: List Core.Statement := [.havoc (PyExprToString lhs) md]
    let else_ss: List Core.Statement := [.set (PyExprToString lhs) (.op () "ListStr_nil" none) md]
    res.stmts ++ [.ite (.det guard) then_ss else_ss md]

partial def PyStmtToCore (jmp_targets: List String) (translation_ctx : TranslationContext) (s : stmt SourceRange) : List Core.Statement × TranslationContext :=
  assert! jmp_targets.length > 0
  let md := sourceRangeToMetaData translation_ctx.filePath s.toAst.ann
  let non_throw : List Core.Statement × Option (String × Lambda.LMonoTy) := match s with
    | .Import _ names =>
      ([.call "import" [.inArg (PyListStrToCore names.val)] md], none)
    | .ImportFrom _ s names i =>
      let n := match s.val with
      | some s => [strToCoreExpr s.val]
      | none => []
      let i := match i.val with
      | some i => [intToCoreExpr (PyIntToInt i)]
      | none => []
      ([.call "importFrom" ((n ++ [PyListStrToCore names.val] ++ i).map .inArg) md], none)
    | .Expr _ (.Call _ func args kwords) =>
      let fname := PyExprToString func
      if callCanThrow translation_ctx.func_infos s then
        (handleFunctionCall ["maybe_except"] fname args kwords jmp_targets translation_ctx s, none)
      else
        (handleFunctionCall [] fname args kwords jmp_targets translation_ctx s, none)
    | .Expr _ (.Constant _ (.ConString _ _) _) =>
      -- TODO: Check that it's a doc string
      ([], none) -- Doc string
    | .Expr _ _ =>
      panic! s!"Can't handle Expr statements that aren't calls: {repr s}"
    | .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 translation_ctx s, none)
    | .Assign _ lhs rhs _ =>
      assert! lhs.val.size == 1
      let res := PyExprToCore translation_ctx rhs
      (res.stmts ++ [.set (PyExprToString lhs.val[0]!) res.expr md], none)
    | .AnnAssign _ lhs ty { ann := _ , val := (.some (.Call _ func args kwords))} _ =>
      let fname := PyExprToString func
      (handleFunctionCall [PyExprToString lhs, "maybe_except"] fname args kwords jmp_targets translation_ctx s, some (PyExprToString lhs, PyExprToMonoTy ty))
    | .AnnAssign _ lhs ty { ann := _ , val := (.some (.ListComp _ _ gen))} _ =>
      (handleComprehension translation_ctx lhs gen.val, some (PyExprToString lhs, PyExprToMonoTy ty))
    | .AnnAssign _ lhs ty {ann := _, val := (.some e)} _ =>
      let res := (PyExprToCore {translation_ctx with expectedType := PyExprToMonoTy ty} e)
      (res.stmts ++ [.set (PyExprToString lhs) res.expr md], some (PyExprToString lhs, PyExprToMonoTy ty))
    | .Try _ body handlers _orelse _finalbody =>
        let new_target := s!"excepthandlers_{jmp_targets[0]!}"
        let new_jmp_stack := new_target :: jmp_targets
        let except_handlers := handlers.val.toList.flatMap (exceptHandlersToCore new_jmp_stack translation_ctx)
        let var_decls := collectVarDecls translation_ctx body.val
        let try_body := var_decls ++ body.val.toList.flatMap (λ s => (PyStmtToCore new_jmp_stack translation_ctx s).fst)
        ([.block "try_block" ([.block new_target try_body .empty] ++ except_handlers) .empty], none)
    | .FunctionDef _ _ _ _ _ _ _ _ => panic! "Can't translate FunctionDef to Strata Core statement"
    | .If _ test then_b else_b =>
      let guard_ctx := {translation_ctx with expectedType := some (.tcons "bool" [])}
      ([.ite (.det (PyExprToCore guard_ctx test).expr) (ArrPyStmtToCore translation_ctx then_b.val).fst (ArrPyStmtToCore translation_ctx else_b.val).fst md], none)
    | .Return _ v =>
      match v.val with
      | .some v => ([.set "ret" (PyExprToCore translation_ctx v).expr md, .exit jmp_targets[0]! md], none) -- TODO: need to thread return value name here. For now, assume "ret"
      | .none => ([.exit jmp_targets[0]! md], none)
    | .For _ tgt itr body _ _ =>
      -- Do one unrolling:
      let guard := .app () (Core.coreOpExpr (.bool .Not)) (.eq () (.app () (.op () "dict_str_any_length" none) (PyExprToCore default itr).expr) (.intConst () 0))
      match tgt with
      | .Name _ n _ =>
        let assign_tgt := [(.init n.val dictStrAnyType (.det dummyDictStrAny) md)]
        ([.ite (.det guard) (assign_tgt ++ (ArrPyStmtToCore translation_ctx body.val).fst) [] md], none)
      | _ => panic! s!"tgt must be single name: {repr tgt}"
      -- TODO: missing havoc
    | .While _ test body _ =>
      -- Do one unrolling:
      let guard := .app () (Core.coreOpExpr (.bool .Not)) (.eq () (.app () (.op () "dict_str_any_length" none) (PyExprToCore default test).expr) (.intConst () 0))
      ([.ite (.det guard) (ArrPyStmtToCore translation_ctx body.val).fst [] md], none)
      -- TODO: missing havoc
    | .Assert sr a _ =>
      let res := PyExprToCore translation_ctx a
      let md := sourceRangeToMetaData translation_ctx.filePath sr
      ([(.assert "py_assertion" res.expr md)], none)
    | .AugAssign _ lhs op rhs =>
      match op with
      | .Add _ =>
        match lhs with
        | .Name _ n _ =>
          let rhs := PyExprToCore translation_ctx rhs
          let new_lhs := (.strConst () "DUMMY_FLOAT")
          (rhs.stmts ++ [.set n.val new_lhs md], none)
        | _ => panic! s!"Expected lhs to be name: {repr lhs}"
      | .FloorDiv _ =>
        match lhs with
        | .Name _ n _ =>
          let lhs := PyExprToCore translation_ctx lhs
          let rhs := PyExprToCore translation_ctx rhs
          let new_lhs := .app () (.app () (Core.coreOpExpr (.numeric ⟨.int, .Div⟩) (some mty[int → (int → int)])) lhs.expr) rhs.expr
          (rhs.stmts ++ [.set n.val new_lhs md], none)
        | _ => panic! s!"Expected lhs to be name: {repr lhs}"
      | _ => panic! s!"Unsupported AugAssign op: {repr op}"
    | _ =>
      panic! s!"Unsupported {repr s}"
  let new_translation_ctx := match non_throw.snd with
  | .some s => {translation_ctx with variableTypes := s :: translation_ctx.variableTypes}
  | .none => translation_ctx
  if callCanThrow translation_ctx.func_infos s then
    (non_throw.fst ++ [handleCallThrow jmp_targets[0]!], new_translation_ctx)
  else
    (non_throw.fst, new_translation_ctx)

partial def ArrPyStmtToCore (translation_ctx: TranslationContext) (a : Array (stmt SourceRange)) : (List Core.Statement × TranslationContext) :=
  a.foldl (fun (stmts, ctx) stmt =>
    let (newStmts, newCtx) := PyStmtToCore ["end"] ctx stmt
    (stmts ++ newStmts, newCtx)
  ) ([], translation_ctx)

end --mutual


def translateFunctions (a : Array (stmt SourceRange)) (translation_ctx: TranslationContext) : List Core.Decl :=
  a.toList.filterMap (λ s => match s with
    | .FunctionDef _ name _args body _ _ret _ _ =>

      let varDecls : List Core.Statement := []
      let proc : Core.Procedure := {
        header := {
               name := name.val,
               typeArgs := [],
               inputs := [],
               outputs := [("maybe_except", (.tcons "ExceptOrNone" []))]},
        spec := default,
        body := .structured (varDecls ++ [.block "end" ((ArrPyStmtToCore translation_ctx body.val).fst) .empty])
      }
      some (.proc proc .empty)
    | _ => none)

def pyTyStrToLMonoTy (ty_str: String) : Lambda.LMonoTy :=
  match ty_str with
  | "str" => mty[string]
  | "int" => mty[int]
  | "datetime" => (.tcons "Datetime" [])
  | _ => panic! s!"Unsupported type: {ty_str}"

def pythonFuncToCore (name : String) (args: List (String × String)) (body: Array (stmt SourceRange)) (ret : Option (expr SourceRange)) (spec : Core.Procedure.Spec) (translation_ctx : TranslationContext) : Core.Procedure :=
  let inputs : List (Lambda.Identifier Unit × Lambda.LMonoTy) := args.map (λ p => (p.fst, pyTyStrToLMonoTy p.snd))
  let varDecls := collectVarDecls translation_ctx body ++ [(.init "exception_ty_matches" t[bool] (.det (.boolConst () false)) .empty), (.havoc "exception_ty_matches" .empty)]
  let stmts := (ArrPyStmtToCore translation_ctx body).fst
  let body := varDecls ++ [.block "end" stmts .empty]
  let constructor := name.endsWith "___init__"
  let outputs : Lambda.LMonoTySignature := if not constructor then
    match ret with
    | .some _v => [("ret", (.tcons "DictStrAny" [])), ("maybe_except", (.tcons "ExceptOrNone" []))]
    | .none => [("maybe_except", (.tcons "ExceptOrNone" []))]
  else
    let class_ty_name := name.dropEnd ("___init__".length) |>.toString
    [("ret", (.tcons s!"{class_ty_name}" [])), ("maybe_except", (.tcons "ExceptOrNone" []))]
  {
    header := {name,
               typeArgs := [],
               inputs,
               outputs},
    spec,
    body := .structured body
  }

def unpackPyArguments (args: arguments SourceRange) : List (String × String) :=
-- Python AST:
-- arguments = (arg* posonlyargs, arg* args, arg? vararg, arg* kwonlyargs,
--                  expr* kw_defaults, arg? kwarg, expr* defaults)
  match args with -- TODO: Error if any other types of args
  | .mk_arguments _ _ args _ _ _ _ _ =>
    let combined := args.val
    combined.toList.filterMap (λ a =>
    match a with
    | .mk_arg _ name oty _ =>
      if name.val == "self" then
        none
      else
        match oty.val with
        | .some ty => some (name.val, PyExprToString ty)
        | _ => panic! s!"Missing type annotation on arg: {repr a} ({repr args})")

def PyFuncDefToCore (s: stmt SourceRange) (translation_ctx: TranslationContext) : List Core.Decl × PythonFunctionDecl :=
  match s with
  | .FunctionDef _ name args body _ ret _ _ =>
    let args := unpackPyArguments args
    ([.proc (pythonFuncToCore name.val args body.val ret.val default translation_ctx) .empty],
     {name := name.val, args, ret := s!"{repr ret}"})
  | _ => panic! s!"Expected function def: {repr s}"

def PyClassDefToCore (s: stmt SourceRange) (translation_ctx: TranslationContext) : List Core.Decl × PythonClassDecl :=
  match s with
  | .ClassDef _ c_name _ _ body _ _ =>
    let member_fn_defs := body.val.toList.filterMap (λ s => match s with
      | .FunctionDef _ name args body _ ret _ _ => some (name, args, body, ret)
      | _ => none)
    (member_fn_defs.map (λ f =>
      let name := f.fst.val
      let args := unpackPyArguments f.snd.fst
      let body := f.snd.snd.fst.val
      let ret := f.snd.snd.snd.val
      .proc (pythonFuncToCore (c_name.val++"_"++name) args body ret default translation_ctx) .empty),
      {name := c_name.val})
  | _ => panic! s!"Expected function def: {repr s}"

def pythonToCore (signatures : Signatures) (insideMod : Array (stmt SourceRange)) (prelude : Core.Program) (filePath : String := ""): Core.Program :=
  let func_defs := insideMod.filter (λ s => match s with
  | .FunctionDef _ _ _ _ _ _ _ _ => true
  | _ => false)

  let class_defs := insideMod.filter (λ s => match s with
  | .ClassDef _ _ _ _ _ _ _ => true
  | _ => false)

  let non_func_blocks := insideMod.filter (λ s => match s with
  | .FunctionDef _ _ _ _ _ _ _ _ => false
  | .ClassDef _ _ _ _ _ _ _ => false
  | _ => true)

  let rec helper {α : Type} (f : stmt SourceRange → TranslationContext → List Core.Decl × α)
               (update : TranslationContext → α → TranslationContext)
               (acc : TranslationContext) :
               List (stmt SourceRange) → List Core.Decl × TranslationContext
  | [] => ([], acc)
  | x :: xs =>
    let (y, info) := f x acc
    let new_acc := update acc info
    let (ys, acc'') := helper f update new_acc xs
    (y ++ ys, acc'')

  -- TODO: in Python, declarations can be circular
  let preludeProcNames := extractProcedureNames prelude
  let base_ctx : TranslationContext := { signatures, filePath, preludeProcedureNames := preludeProcNames }

  let class_defs_and_infos := helper PyClassDefToCore (fun acc info => {acc with class_infos := info :: acc.class_infos}) base_ctx class_defs.toList
  let class_defs := class_defs_and_infos.fst
  let class_infos := class_defs_and_infos.snd

  let class_ty_decls := class_infos.class_infos.map (λ info =>
    .type (.con {name := info.name, params := []}) .empty)

  let func_defs_and_infos := helper PyFuncDefToCore (fun acc info => {acc with func_infos := info :: acc.func_infos}) class_infos func_defs.toList
  let func_defs := func_defs_and_infos.fst
  let func_infos := func_defs_and_infos.snd

  {decls := class_ty_decls ++ func_defs ++ class_defs ++
    [.proc (pythonFuncToCore "__main__" [] non_func_blocks none default func_infos) .empty]}

end -- public section
end StrataPython