Strata/Strata/Languages/Python/FunctionSignatures.lean at e4723b66ff9b76d349fbdd4fbfa949f87ed3ee49 · strata-org/Strata · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
/-
  Copyright Strata Contributors

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

import Strata.Languages.Boogie.Boogie

namespace Strata
namespace Python

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
    panic! s!"Should only be called for possibly None types. Called for: {ty}"
  else
    match ty with
    | "StrOrNone" => .app () (.op () "StrOrNone_mk_none" none) (.op () "None_none" none)
    | "BoolOrNone" => .app () (.op () "BoolOrNone_mk_none" none) (.op () "None_none" none)
    | "BoolOrStrOrNone" => .app () (.op () "BoolOrStrOrNone_mk_none" none) (.op () "None_none" none)
    | "AnyOrNone" => .app () (.op () "AnyOrNone_mk_none" none) (.op () "None_none" none)
    | "IntOrNone" => .app () (.op () "IntOrNone_mk_none" none) (.op () "None_none" none)
    | "BytesOrStrOrNone" => .app () (.op () "BytesOrStrOrNone_mk_none" none) (.op () "None_none" none)
    | "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