PySpecArgTypeTest.lean

/-
  Copyright Strata Contributors

  SPDX-License-Identifier: Apache-2.0 OR MIT
-/
module

public import StrataPython.PythonIdent
meta import all StrataPython.PySpecPipeline
meta import all StrataPython.Specs.DDM
import Strata.Languages.Laurel.Grammar.AbstractToConcreteTreeTranslator

meta section

/-! ## Test: specArgToFuncDeclArg preserves parameter type info

Verifies that `buildPySpecLaurel` populates `laurelType` on `PyArgInfo`
and that `typeTesters` is empty (type assertions live in the procedure body).
-/

namespace StrataPython.PySpecArgTypeTest

open StrataPython.Specs
open StrataPython (buildPySpecLaurel)
open StrataPython (ModuleName OverloadTable)
open StrataPython.ToLaurel (PythonFunctionDecl PyArgInfo highTypeToPyLauType)
open Strata
open Strata.Laurel (Procedure formatProcedure)

private def loc : SourceRange := default

private def identType (nm : PythonIdent) : SpecType :=
  SpecType.ident loc nm

private def mkArg (name : String) (type : SpecType) : Specs.Arg :=
  { name, type }

private def mkFunc (name : String) (args : Array Specs.Arg) (ret : SpecType) : Signature :=
  .functionDecl {
    loc, nameLoc := loc, name
    args := { args, kwonly := #[] }
    returnType := ret
    isOverload := false
    preconditions := #[], postconditions := #[]
  }

/-- Build PySpec signatures, write to temp Ion, read back via buildPySpecLaurel,
    and return the full result. -/
private def buildSpecs (sigs : Array Signature) : IO StrataPython.PySpecLaurelResult := do
  IO.FS.withTempDir fun dir => do
    let ionFile := dir / "test.pyspec.ion"
    writeDDM ionFile sigs
    let ctx ← Strata.Pipeline.PipelineContext.create
    match ← (buildPySpecLaurel ctx #[(.ofComponent (.ofString "test"), ionFile.toString)] {}).toBaseIO with
    | .ok r => pure r
    | .error () =>
      let msgs ← ctx.getMessages
      throw <| .userError s!"buildPySpecLaurel failed: {msgs.map toString}"

private def getFuncSigs (sigs : Array Signature) : IO (List PythonFunctionDecl) := do
  return (← buildSpecs sigs).functionSignatures

private def unionType (elts : Array SpecType) : SpecType :=
  SpecType.unionArray loc elts

/--
info: test_typed_func: x=Any[], y=Any[], z=Any[], w=Any[]
test_untyped_func: a=Any[]
test_mixed_func: p=Any[], q=Any[]
test_optional_func: s=Any[], n=Any[]
-/
#guard_msgs in
#eval do
  let sigs ← getFuncSigs #[
    mkFunc "typed_func"
      #[mkArg "x" (identType .builtinsInt),
        mkArg "y" (identType .builtinsStr),
        mkArg "z" (identType .builtinsBool),
        mkArg "w" (identType .builtinsFloat)]
      (identType .noneType),
    mkFunc "untyped_func"
      #[mkArg "a" (identType .typingAny)]
      (identType .noneType),
    mkFunc "mixed_func"
      #[mkArg "p" (identType .builtinsStr),
        mkArg "q" (identType .typingAny)]
      (identType .noneType),
    mkFunc "optional_func"
      #[mkArg "s" (unionType #[identType .noneType, identType .builtinsStr]),
        mkArg "n" (unionType #[identType .noneType, identType .builtinsInt])]
      (identType .noneType)
  ]
  for f in sigs do
    let argStrs := ", ".intercalate (f.args.map fun (a : PyArgInfo) =>
      s!"{a.name}={highTypeToPyLauType a.laurelType.val}{a.typeTesters.toList}")
    IO.println s!"{f.name}: {argStrs}"

/-! ## Test: buildSpecBody generates type assertions in procedure body

Verifies that even though `typeTesters` on `PythonFunctionDecl` is empty,
the pyspec Laurel procedure body contains the type assertions generated by
`buildSpecBody`. This confirms the body-side checks make call-site
preconditions redundant. -/

/--
info: procedure test_typed_func(x: Any, y: Any): Any
  opaque
  modifies *
{
  result := <??>;
  assert Any..isfrom_int(x);
  assert Any..isfrom_str(y);
  assume Any..isfrom_float(result)
};
-/
#guard_msgs in
#eval! do
  let result ← buildSpecs #[
    mkFunc "typed_func"
      #[mkArg "x" (identType .builtinsInt),
        mkArg "y" (identType .builtinsStr)]
      (identType .builtinsFloat)
  ]
  let procs := result.laurelProgram.staticProcedures
  let some proc := procs.find? (fun (p : Procedure) => p.name.text == "test_typed_func")
    | throw <| IO.userError "test_typed_func not found"
  IO.println (toString (formatProcedure proc))

end StrataPython.PySpecArgTypeTest
end