fix: handle arguments and defaults consistently

Default function arguments were not being handled the same way as
function arguments. It was possible that an expression that could be
passed as an argument could not have been used as a default value for
the same argument. This patch fixes this issue, and simplifies some of
the parsing code for arguments and function signatures.

NKL/Python.lean

@@ -115,10 +115,10 @@ missing arguments, but this is just how it works in the python AST.
 structure Args where
   posonlyargs : List String
   args : List String
-  defaults: List Expr
+  defaults: List Expr'
   vararg : Option String
   kwonlyargs : List String
-  kw_defaults: List Expr
+  kw_defaults: List (String × Expr')
   kwarg : Option String
   deriving Repr
 
@@ -129,16 +129,14 @@ def Args.names (ax : Args) : List String :=
   let xs := match ax.kwarg  with | none => xs | some x => xs.append [x]
   xs
 
-/-
-In addition to the defaults above from the AST, we also collect
-the values from f.__defaults__ here in the Fun structure. These
-values are evaluated in a different context from the other names
-in the function, so we need to capture them on the Python side.
--/
+def Args.all_defaults (ax : Args) : List (String × Expr') :=
+  let args := ax.posonlyargs ++ ax.args
+  let dflt := args.reverse.zip ax.defaults.reverse
+  dflt ++ ax.kw_defaults
+
 structure Fun where
   source : String
   args : Args
-  defaults: List Const
   body: List Stmt
   deriving Repr
 
@@ -315,33 +313,6 @@ where
     | "If" => return (.ifStm (<- expr "test") (<- stmts "body") (<- stmts "orelse"))
     | _ => throw s!"unsupported python construct {key}"
 
-def arguments (j : Json) : Parser Args := do
-  let obj <- j.getObjVal? "arguments"
-  let arg? := field (opt arg) obj
-  let args := field (list arg) obj
-  let exprs := field (list expr) obj
-  return {
-    posonlyargs := (<- args "posonlyargs")
-    args        := (<- args "args")
-    defaults    := (<- exprs "defaults")
-    vararg      := (<- arg? "vararg")
-    kwonlyargs  := (<- args "kwonlyargs")
-    kw_defaults := (<- exprs "kw_defaults")
-    kwarg       := (<- arg? "kwarg")
-  }
-where
-  arg (j : Json) : Parser String := do
-    let obj <- j.getObjVal? "arg"
-    return (<- field str obj "arg")
-
-def function (j : Json) : Parser Fun := do
-  let source <- field str j "source"
-  withSrc source do
-    let args <- field arguments j "args"
-    let defaults <- field (list const) j "defaults"
-    let body <- field (list stmt) j "body"
-    return Fun.mk source args defaults body
-
 -- Both global references and arguments are processed in the global
 -- environment. These terms do not have a position, and must be
 -- evaluable in the default environment.
@@ -366,6 +337,29 @@ where
   globals (arr : Array Json) : Parser (List Expr) :=
     arr.toList.mapM fun x => return .exprPos (<- global x) {}
 
+def arguments (j : Json) : Parser Args := do
+  let obj <- j.getObjVal? "arguments"
+  return {
+    posonlyargs := (<- field (list arg)    obj "posonlyargs")
+    args        := (<- field (list arg)    obj "args")
+    defaults    := (<- field (list global) obj "defaults")
+    vararg      := (<- field (opt arg)     obj "vararg")
+    kwonlyargs  := (<- field (list arg)    obj "kwonlyargs")
+    kw_defaults := (<- field (dict global) obj "kw_defaults")
+    kwarg       := (<- field (opt arg)     obj "kwarg")
+  }
+where
+  arg (j : Json) : Parser String := do
+    let obj <- j.getObjVal? "arg"
+    return (<- field str obj "arg")
+
+def function (j : Json) : Parser Fun := do
+  let source <- field str j "source"
+  withSrc source do
+    let args <- field arguments j "args"
+    let body <- field (list stmt) j "body"
+    return Fun.mk source args body
+
 def kernel (j : Json) : Parser Kernel := do
   let name <- field str j "entry"
   let funcs <- field (dict function) j "funcs"

interop/nkl/__init__.py

@@ -2,5 +2,4 @@
 # Released under Apache 2.0 license as described in the file LICENSE.
 # Authors: Paul Govereau
 
-from .lean import load, to_json
 from .parser import Parser

interop/nkl/parser.py

@@ -11,7 +11,7 @@
 from textwrap import dedent
 from itertools import chain
 from collections import deque
-from nkl.lean import py_to_lean
+from nkl.lean_rffi import py_to_lean
 
 # This is a custom JSON encoder for use with AST nodes.
 # The AST nodes are not handled by the default encoder.
@@ -97,25 +97,29 @@ def json(self):
   def load(self):
     py_to_lean(self.json())
 
-  def apply_args(self, *args, **kwargs):
-    self.args = []
-    self.kwargs = {}
+  def process_args(self, args, kwargs):
+    l = []
     d = {}
-    for arg in args:
-      self.reference(d, '_', arg)
-      try: self.args.append(d.popitem()[1])
-      except Exception:
-        raise Exception("Unsupported argument type")
-    for k,v in kwargs.items():
-      self.ref_arg(k, v)
+    if args:
+      for arg in args:
+        self.reference(d, '_', arg)
+        try: l.append(d.popitem()[1])
+        except Exception:
+          raise Exception("Unsupported argument type")
+    if kwargs:
+      for k,v in kwargs.items():
+        self.reference(d, k, v)
+    return l, d
+
+  def apply_args(self, *args, **kwargs):
+    l, d = self.process_args(args, kwargs)
+    self.args = l
+    self.kwargs = d
 
   def __call__(self, *args, **kwargs):
     self.apply_args(*args, **kwargs)
     py_to_lean(self.json())
 
-  def ref_arg(self, refname, val):
-    return self.reference(self.kwargs, refname, val)
-
   def ref_global(self, refname, val):
     return self.reference(self.globals, refname, val)
 
@@ -160,9 +164,11 @@ def translate(self, f: types.FunctionType, args: ast.arguments, body: [ast.AST])
     self.f = f
     for s in body:
       self.visit(s)
+    l, d = self.process_args(f.__defaults__, f.__kwdefaults__)
+    args.defaults = l
+    args.kw_defaults = d
     return { 'source': inspect.getsource(f)
            , 'args': args
-           , 'defaults': list(self.fun_defaults(f))
            , 'body': body
            }
 
@@ -201,20 +207,3 @@ def visit_Attribute(self, node):
       raise e
     except Exception:
       return
-
-  def fun_defaults(self, f: types.FunctionType):
-    if f.__defaults__ is None:
-      return dict()
-    names = f.__code__.co_varnames[:f.__code__.co_argcount]
-    tbl = { n:v for (n,v) in zip(reversed(names), reversed(f.__defaults__)) }
-    if f.__kwdefaults__ is not None:
-      tbl.update(f.__kwdefaults__)
-    def is_ok(x):
-      if x is None or isinstance(x, (int, float, str)):
-        return True
-      if isinstance(x, types.FunctionType):
-        # TODO: this could be incorrect if default
-        # is using an alternate name for the function
-        self.ref_global(x.__name__, x)
-      return False
-    return { n:v for (n,v) in tbl.items() if is_ok(v) }