packing and unpacking arguments

Stream: Is there code for X?

[Jianlin Li]

Hi everyone! :wave:

Does Lean have anything like Python’s **kwargs (keyword-argument unpacking)?

Lean lets me call functions with named arguments out of order, e.g.:

def f (n : Nat) (m : Nat) : Nat :=   n * m
#eval f 3 4                 -- 12
#eval f (n := 5) (m := 6)   -- 30
#eval f (m := 6) (n := 5)   -- 30
#check f (m := 10)          -- fun n => f n 10 : Nat → Nat

What I’m looking for is the Python-style behavior where a runtime dictionary can be unpacked into named arguments:

def f(n=0, m=0): return n*m
kwargs = {"n": 3, "m": 4}
f(**kwargs)  # 12

Is there a way in Lean to pass a runtime map (e.g., HashMap String Nat) and have it supply named parameters to a function? If not built-in, are there libraries, patterns, or metaprogramming approaches to simulate this (possibly via records/structures, HashMaps, currying/uncurrying, or macros)?

[Michael MacLeod]

Here's one option. It's a bit of a hack, but I think it does what you want. It uses default parameters and looks them up from an args hashmap if they are zero:

def f (n : Nat := 0) (m : Nat := 0) (args : Std.HashMap String Nat := .emptyWithCapacity 0) :=
  let n := args.get? "n" |>.getD n
  let m := args.get? "m" |>.getD m
  n*m

#eval f -- 0
#eval f 2 3 -- 6
#eval f (args := .ofList [("n", 3), ("m", 4)]) -- 12

I'm not too sure what the use of doing this is though. Maybe default parameters are all you need?

[Kim Morrison]

I think @Jianlin Li is asking about a macro that would make this possible for any function, not how to implement a single function this way.

[Siddhartha Gadgil]

Should be possible with an elaborator matching the type of f as a pi-type. Will give it a shot if nobody else does it first.

An issue: what is the type of the HashMap? We need something polymorphic unless the arguments all have the same type. Maybe some form of HList should be used.

[Siddhartha Gadgil]

And a search for HList brings up Mario's suggestion to use Prod as HList which fits well here.

[Siddhartha Gadgil]

Here is the start of an implementation.

import Lean

open Lean Meta Elab Term

def isProd? (e: Expr) : MetaM (Option <| Expr × Expr) := do
  let u ← mkFreshLevelMVar
  let v ← mkFreshLevelMVar
  let u := mkSort <| Level.succ u
  let v := mkSort <| Level.succ v
  let α ← mkFreshExprMVar u
  let β ← mkFreshExprMVar v
  let x ←  mkFreshExprMVar α
  let y ←  mkFreshExprMVar β
  let p ← mkAppM ``Prod.mk #[x, y]
  if ← isDefEq p e then
    return some (x, y)
  return none

partial def readKeyVal (e: Expr)  : MetaM (Option <| Name × Expr) := do
  match ← isProd? e with
  | some (kExpr, v) => do
    let α ← inferType kExpr
    unless ← isDefEq α (mkConst ``Name) do
      return none
    let key ← unsafe evalExpr Name (mkConst ``Name) kExpr
    return (key, v)
  | none => do
    return none

partial def readKVs (e: Expr) : MetaM (Std.HashMap Name Expr) := do
  match ←  readKeyVal e with
  | some (key, value) => do
    let mut result := Std.HashMap.empty
    result := result.insert key value
    return result
  | none => do
    match ← isProd? e with
    | some (head, tail) => do
      match ← readKeyVal head with
      | some (key, value) => do
        let mut result ← readKVs tail
        result := result.insert key value
        return result
      | none => do
        return Std.HashMap.empty
    | none => do
      return Std.HashMap.empty

-- from batteries
def getExplicitArgsFromType : Expr → Array Name → Array Name
  | .forallE n _ body bi, args =>
    getExplicitArgsFromType body <| if bi.isExplicit then args.push n else args
  | _, args => args

def fillFuncArgs (f: Expr) (kwArgs : Std.HashMap Name Expr) : TermElabM Expr := do
  let mut args : Array Expr := #[]
  let explNames := getExplicitArgsFromType (← inferType f) #[]
  for n in explNames do
    match kwArgs.get? n with
    | some v => args := args.push v
    | none => return f
  Term.synthesizeSyntheticMVarsNoPostponing
  mkAppM' f args

elab f:term "(**" kw:term ")" : term => do
  let f ←
    withoutPostponing do elabTerm f none
  let kw ←
    withoutPostponing do elabTerm kw none
  let kwArgs ← readKVs kw
  fillFuncArgs f kwArgs

def kwargs := ((`a, 1), (`b, 2), (`x, "hello"))

def f (a b :Nat) := a + b
#eval f (** ((`a, 1), (`b, 2), (`x, "hello"))) -- 3
#eval f (** kwargs) -- 3

[Siddhartha Gadgil]

This is not tested much. One may also want nicer syntax, but a macro should fix that.

[Siddhartha Gadgil]

Here is the code with some syntax added:

import Lean

open Lean Meta Elab Command Syntax Term Parser

def isProd? (e: Expr) : MetaM (Option <| Expr × Expr) := do
  let u ← mkFreshLevelMVar
  let v ← mkFreshLevelMVar
  let u := mkSort <| Level.succ u
  let v := mkSort <| Level.succ v
  let α ← mkFreshExprMVar u
  let β ← mkFreshExprMVar v
  let x ←  mkFreshExprMVar α
  let y ←  mkFreshExprMVar β
  let p ← mkAppM ``Prod.mk #[x, y]
  if ← isDefEq p e then
    return some (x, y)
  return none

partial def kv? (e: Expr)  : MetaM (Option <| Name × Expr) := do
  match ← isProd? e with
  | some (kExpr, v) => do
    let α ← inferType kExpr
    unless ← isDefEq α (mkConst ``Name) do
      return none
    let key ← unsafe evalExpr Name (mkConst ``Name) kExpr
    return (key, v)
  | none => do
    return none

partial def kvs? (e: Expr) : MetaM (Std.HashMap Name Expr) := do
  match ←  kv? e with
  | some (key, value) => do
    let mut result := Std.HashMap.empty
    result := result.insert key value
    return result
  | none => do
    match ← isProd? e with
    | some (head, tail) => do
      match ← kv? head with
      | some (key, value) => do
        let mut result ← kvs? tail
        result := result.insert key value
        return result
      | none => do
        return Std.HashMap.empty
    | none => do
      return Std.HashMap.empty

-- from batteries
def getExplicitArgsFromType : Expr → Array Name → Array Name
  | .forallE n _ body bi, args =>
    getExplicitArgsFromType body <| if bi.isExplicit then args.push n else args
  | _, args => args

def fillFuncKWArgs (f: Expr) (kwArgs : Std.HashMap Name Expr) : TermElabM Expr := do
  let mut args : Array Expr := #[]
  let explNames := getExplicitArgsFromType (← inferType f) #[]
  for n in explNames do
    match kwArgs.get? n with
    | some v => args := args.push v
    | none => return f
  Term.synthesizeSyntheticMVarsNoPostponing
  mkAppM' f args

elab f:term "(**" kw:term ")" : term => do
  let f ←
    withoutPostponing do elabTerm f none
  let kw ←
    withoutPostponing do elabTerm kw none
  let kwArgs ← kvs? kw
  fillFuncKWArgs f kwArgs

def kwargsEg := ((`a, 1), (`b, 2), (`x, "hello"))

def f (a b :Nat) := a + b
#eval f (** ((`a, 1), (`b, 2), (`x, "hello"))) -- 3
#eval f (** kwargsEg) -- 3


syntax assgn := ident " := " term
syntax (name:=kwargs) "{{" assgn,* "}}" : term

def kwTerm : TSyntax ``assgn → MacroM Syntax.Term
  | `(assgn| $n:ident := $v:term) => do
    let n := quoteNameMk n.getId
    `(($n, $v))
  | _ => throw Lean.Macro.Exception.unsupportedSyntax

macro_rules
| `({{ $p:assgn }}) =>
    kwTerm p
| `({{ $as:assgn,*, $last:assgn }}) => do
    let head ← kwTerm last
    let tailPairs ← Array.mapM kwTerm as
    let stx ← tailPairs.foldrM (init := head) (fun acc p => `(($acc, $p)))
    return stx

#eval f (** kwargsEg) -- 3

def kwArgsEg' := {{ a := 1, b := 2, c := "hello" }}

#eval kwargsEg
#eval kwArgsEg'

#eval f (** kwArgsEg')
#eval f (** {{ a := 1, b := 2, c := "hello" }}) -- 3

[Siddhartha Gadgil]

The doubled braces were to avoid collisions with structures but Lean seems smart enough that they are not needed.

[Adam Topaz]

Siddhartha Gadgil said:

Should be possible with an elaborator matching the type of f as a pi-type. Will give it a shot if nobody else does it first.

An issue: what is the type of the HashMap? We need something polymorphic unless the arguments all have the same type. Maybe some form of HList should be used.

We do have docs#Std.DHashMap