Add polymorphic functions to DDM and Strata.Boogie

Strata/DDM/AST.lean

@@ -90,6 +90,9 @@ inductive TypeExprF (α : Type) where
 | ident (ann : α) (name : QualifiedIdent) (args : Array (TypeExprF α))
   /-- A bound type variable at the given deBruijn index in the context. -/
 | bvar (ann : α) (index : Nat)
+  /-- A polymorphic type variable (universally quantified).
+      Used for polymorphic function type parameters -/
+| tvar (ann : α) (name : String)
   /-- A reference to a global variable along with any arguments to ensure it is well-typed. -/
 | fvar (ann : α) (fvar : FreeVarIndex) (args : Array (TypeExprF α))
   /-- A function type. -/
@@ -101,6 +104,7 @@ namespace TypeExprF
 def ann {α} : TypeExprF α → α
 | .ident ann _ _ => ann
 | .bvar ann _ => ann
+| .tvar ann _ => ann
 | .fvar ann _ _ => ann
 | .arrow ann _ _ => ann
 
@@ -112,13 +116,15 @@ protected def incIndices {α} (tp : TypeExprF α) (count : Nat) : TypeExprF α :
   | .ident n i args => .ident n i (args.attach.map fun ⟨e, _⟩ => e.incIndices count)
   | .fvar n f args => .fvar n f (args.attach.map fun ⟨e, _⟩ => e.incIndices count)
   | .bvar n idx => .bvar n (idx + count)
+  | .tvar n name => .tvar n name  -- tvar doesn't use indices
   | .arrow n a r => .arrow n (a.incIndices count) (r.incIndices count)
 
 /-- Return true if type expression has a bound variable. -/
 protected def hasUnboundVar {α} (bindingCount : Nat := 0) : TypeExprF α → Bool
 | .ident _ _ args => args.attach.any (fun ⟨e, _⟩ => e.hasUnboundVar bindingCount)
 | .fvar _ _ args => args.attach.any (fun ⟨e, _⟩ => e.hasUnboundVar bindingCount)
 | .bvar _ idx => idx ≥ bindingCount
+| .tvar _ _ => true
 | .arrow _ a r => a.hasUnboundVar bindingCount || r.hasUnboundVar bindingCount
 termination_by e => e
 
@@ -135,6 +141,7 @@ protected def instTypeM {m α} [Monad m] (d : TypeExprF α) (bindings : α → N
   | .ident n i a =>
     .ident n i <$> a.attach.mapM (fun ⟨e, _⟩ => e.instTypeM bindings)
   | .bvar n idx => bindings n idx
+  | .tvar n name => pure (.tvar n name)
   | .fvar n idx a => .fvar n idx <$> a.attach.mapM (fun ⟨e, _⟩ => e.instTypeM bindings)
   | .arrow n a b => .arrow n <$> a.instTypeM bindings <*> b.instTypeM bindings
 termination_by d
@@ -472,6 +479,15 @@ def scopeDatatypeIndex (metadata : Metadata) : Option (Nat × Nat) :=
   | some #[.catbvar nameIdx, .catbvar typeParamsIdx] => some (nameIdx, typeParamsIdx)
   | some _ => panic! s!"Unexpected argument count to scopeDatatype"
 
+/-- Returns the typeArgs index if @[scopeTypeVars] is present.
+    Used for polymorphic function declarations where type parameters should be
+    added as .tvar bindings. -/
+def scopeTypeVarsIndex (metadata : Metadata) : Option Nat :=
+  match metadata[q`StrataDDL.scopeTypeVars]? with
+  | none => none
+  | some #[.catbvar typeArgsIdx] => some typeArgsIdx
+  | some _ => panic! s!"Unexpected argument count to scopeTypeVars"
+
 /-- Returns the index of the value in the binding for the given variable of the scope to use. -/
 private def resultIndex (metadata : Metadata) : Option Nat :=
   match metadata[MetadataAttr.scopeName]? with
@@ -510,8 +526,12 @@ generate types.
 inductive PreType where
   /-- A dialect defined type. -/
 | ident (ann : SourceRange) (name : QualifiedIdent) (args : Array PreType)
-  /-- A bound type variable at the given deBruijn index in the context. -/
+  /-- A bound type variable at the given deBruijn index in the context.
+      Used for type alias parameters -/
 | bvar (ann : SourceRange) (index : Nat)
+  /-- A polymorphic type variable (universally quantified).
+      Used for polymorphic function type parameters -/
+| tvar (ann : SourceRange) (name : String)
   /-- A reference to a global variable along with any arguments to ensure it is well-typed. -/
 | fvar (ann : SourceRange) (fvar : FreeVarIndex) (args : Array PreType)
   /-- A function type. -/
@@ -526,13 +546,15 @@ namespace PreType
 def ann : PreType → SourceRange
 | .ident ann _ _ => ann
 | .bvar ann _ => ann
+| .tvar ann _ => ann
 | .fvar ann _ _ => ann
 | .arrow ann _ _ => ann
 | .funMacro ann _ _ => ann
 
 def ofType : TypeExprF SourceRange → PreType
 | .ident loc name args => .ident loc name (args.map fun a => .ofType a)
 | .bvar loc idx => .bvar loc idx
+| .tvar loc name => .tvar loc name
 | .fvar loc idx args => .fvar loc idx (args.map fun a => .ofType a)
 | .arrow loc a r => .arrow loc (.ofType a) (.ofType r)
 termination_by tp => tp
@@ -696,6 +718,18 @@ def argScopeDatatypeLevel (argDecls : ArgDecls) (level : Fin argDecls.size) : Op
     else
       panic! s!"scopeDatatype name index {nameIdx} out of bounds ({level.val})"
 
+/-- Returns the typeArgs level if @[scopeTypeVars] is present.
+    This is used for polymorphic function declarations where type parameters like `<a, b>`
+    should be added as .tvar bindings when parsing parameter types and return type. -/
+def argScopeTypeVarsLevel (argDecls : ArgDecls) (level : Fin argDecls.size) : Option (Fin level.val) :=
+  match argDecls[level].metadata.scopeTypeVarsIndex with
+  | none => none
+  | some typeArgsIdx =>
+    if h : typeArgsIdx < level.val then
+      some ⟨level.val - (typeArgsIdx + 1), by omega⟩
+    else
+      panic! s!"scopeTypeVars typeArgs index {typeArgsIdx} out of bounds ({level.val})"
+
 end ArgDecls
 
 /--
@@ -752,10 +786,10 @@ structure ConstructorInfo where
 
 /--
 Build a TypeExpr reference to the datatype with type parameters, using
-`.fvar` for the datatype's GlobalContext index.
+`.fvar` for the datatype's GlobalContext index and `.tvar` for type parameters.
 -/
 def mkDatatypeTypeRef (ann : SourceRange) (datatypeIndex : FreeVarIndex) (typeParams : Array String) : TypeExpr :=
-  let typeArgs := typeParams.mapIdx fun i _ => TypeExprF.bvar ann i
+  let typeArgs := typeParams.map fun name => TypeExprF.tvar ann name
   TypeExprF.fvar ann datatypeIndex typeArgs
 
 /--

Strata/DDM/BuiltinDialects/StrataDDL.lean

@@ -171,6 +171,8 @@ def StrataDDL : Dialect := BuiltinM.create! "StrataDDL" #[initDialect] do
    used for recursive datatype definitions where the datatype name must be visible when parsing constructor field types (e.g., `tail: List` in
    `Cons(head: int, tail: List)`) -/
   declareMetadata { name := "scopeDatatype", args := #[.mk "name" .ident, .mk "typeParams" .ident] }
+  -- Metadata for bringing type variable names into scope as .tvar bindings.
+  declareMetadata { name := "scopeTypeVars", args := #[.mk "typeArgs" .ident] }
 
 end Strata
 end

Strata/DDM/Elab/Core.lean

@@ -54,6 +54,7 @@ partial def expandMacros (m : DialectMap) (f : PreType) (args : Nat → Option A
   | .arrow loc a b => .arrow loc <$> expandMacros m a args <*> expandMacros m b args
   | .fvar loc i a => .fvar loc i <$> a.mapM fun e => expandMacros m e args
   | .bvar loc idx => pure (.bvar loc idx)
+  | .tvar loc name => pure (.tvar loc name)
   | .funMacro loc i r => do
     let r ← expandMacros m r args
     match args i with
@@ -77,7 +78,7 @@ the head is in a normal form.
 -/
 partial def hnf (tctx : TypingContext) (e : TypeExpr) : TypeExpr :=
   match e with
-  | .arrow .. | .ident .. => e
+  | .arrow .. | .ident .. | .tvar .. => e
   | .fvar _ idx args =>
     let gctx := tctx.globalContext
     match gctx.kindOf! idx with
@@ -94,6 +95,7 @@ partial def hnf (tctx : TypingContext) (e : TypeExpr) : TypeExpr :=
       assert! d.isGround
       hnf (tctx.drop (idx + 1)) d
     | .type _ _ none => e
+    | .tvar _ _ => e
     | _ => panic! "Expected a type"
 
 /--
@@ -176,6 +178,9 @@ def resolveTypeBinding (tctx : TypingContext) (loc : SourceRange) (name : String
     if let .type loc [] _ := k then
       let info : TypeInfo := { inputCtx := tctx, loc := loc, typeExpr := .bvar loc idx, isInferred := false }
       return .node (.ofTypeInfo info) #[]
+    else if let .tvar loc tvarName := k then
+      let info : TypeInfo := { inputCtx := tctx, loc := loc, typeExpr := .tvar loc tvarName, isInferred := false }
+      return .node (.ofTypeInfo info) #[]
     else
       logErrorMF loc mf!"Expected a type instead of {k}"
       return default
@@ -326,7 +331,7 @@ N.B. This expects that macros have already been expanded in e.
 -/
 partial def headExpandTypeAlias (gctx : GlobalContext) (e : TypeExpr) : TypeExpr :=
   match e with
-  | .arrow .. | .ident .. | .bvar .. => e
+  | .arrow .. | .ident .. | .bvar .. | .tvar .. => e
   | .fvar _ idx args =>
     match gctx.kindOf! idx with
     | .expr _ => panic! "Type free variable bound to expression."
@@ -347,6 +352,10 @@ partial def checkExpressionType (tctx : TypingContext) (itype rtype : TypeExpr)
   let itype := tctx.hnf itype
   let rtype := tctx.hnf rtype
   match itype, rtype with
+  -- tvar matches anything (dialect responsible for typechecking)
+  | .tvar _ n1, .tvar _ n2 => return n1 == n2
+  | .tvar _ _, _ => return true
+  | _, .tvar _ _ => return true
   | .ident _ iq ia, .ident _ rq ra =>
     if p : iq = rq ∧ ia.size = ra.size then do
       for i in Fin.range ia.size do
@@ -468,9 +477,12 @@ partial def unifyTypes
       if !(← checkExpressionType tctx inferredType info.typeExpr) then
         logErrorMF exprLoc mf!"Expression has type {withBindings tctx.bindings (mformat inferredType)} when {withBindings tctx.bindings (mformat info.typeExpr)} expected."
       pure args
+  | .tvar _ _ =>
+    -- tvar nodes are passed through without attempting unification
+    pure args
   | .arrow _ ea er =>
     match inferredType with
-    | .ident .. | .bvar .. | .fvar .. =>
+    | .ident .. | .bvar .. | .fvar .. | .tvar .. =>
       logErrorMF exprLoc mf!"Expected {expectedType} when {inferredType} found"
       pure args
     | .arrow _ ia ir =>
@@ -752,6 +764,12 @@ def translateBindingKind (tree : Tree) : ElabM BindingKind := do
             else
               logErrorMF nameLoc mf!"Unexpected arguments to type variable {name}."
               return default
+          else if let .tvar loc tvarName := k then
+            if tpArgs.isEmpty then
+              return .expr (.tvar loc tvarName)
+            else
+              logErrorMF nameLoc mf!"Unexpected arguments to type variable {name}."
+              return default
           else
             logErrorMF nameLoc mf!"Expected a type instead of {k}"
             return default
@@ -795,12 +813,15 @@ def evalBindingSpec
     let (bindings, success) ← runChecked <| elabArgIndex initSize args b.argsIndex fun argLoc b =>
           match b.kind with
           | .expr tp =>
-            pure (b.ident, tp)
+            pure (some (b.ident, tp))
+          | .tvar .. =>
+            pure none
           | .type .. | .cat _ => do
             logError argLoc "Expecting expressions in variable binding"
-            pure default
+            pure none
     if !success then
       return default
+    let bindings := bindings.filterMap id
     let typeTree := args[b.typeIndex.toLevel]
     let kind ←
           match typeTree.info with
@@ -827,6 +848,8 @@ def evalBindingSpec
           match b.kind with
           | .type _ [] _ =>
             pure ()
+          | .tvar _ _ =>
+            pure ()
           | .type .. | .expr _ | .cat _ => do
             logError argLoc s!"{b.ident} must be have type Type instead of {repr b.kind}."
           return b.ident
@@ -955,6 +978,33 @@ def unwrapTree (tree : Tree) (unwrap : Bool) : Arg :=
     | .ofBytesInfo info => .bytes info.loc info.val
     | _ => tree.arg  -- Fallback for non-unwrappable types
 
+/--
+Extract type variable names from a TypeArgs tree (or Option TypeArgs).
+TypeArgs is an operation `type_args` with one child that is `CommaSepBy Ident`.
+Option TypeArgs wraps this in an OptionInfo node.
+-/
+def extractTypeVarNames (tree : Tree) : Array String :=
+  let innerTree :=
+    match tree.info with
+    | .ofOptionInfo _ =>
+      match tree.children[0]? with
+      | none => none
+      | some t => some t
+    | _ => some tree
+  match innerTree with
+  | none => #[]
+  | some typeArgsTree =>
+    -- typeArgsTree should be the type_args operation
+    -- Its first child should be CommaSepBy Ident
+    match typeArgsTree.children[0]? with
+    | none => #[]
+    | some commaSepTree =>
+      -- commaSepTree has CommaSepInfo with children that are Ident trees
+      commaSepTree.children.filterMap fun identTree =>
+        match identTree.info with
+        | .ofIdentInfo info => some info.val
+        | _ => none
+
 mutual
 
 partial def elabOperation (tctx : TypingContext) (stx : Syntax) : ElabM Tree := do
@@ -1017,21 +1067,43 @@ partial def runSyntaxElaborator
             | .ofIdentInfo info => info.val
             | _ => panic! "Expected identifier for datatype name"
           let baseCtx := typeParamsT.resultContext
+          -- Extract type parameter names from the bindings
+          let typeParamNames := baseCtx.bindings.toArray.filterMap fun b =>
+            match b.kind with
+            | .type _ [] _ => some b.ident
+            | _ => none
           -- Add the datatype name to the GlobalContext as a type
           let gctx := baseCtx.globalContext
           let gctx :=
             if datatypeName ∈ gctx then gctx
-            else gctx.push datatypeName (GlobalKind.type [] none)
-          -- Create a new typing context with the updated GlobalContext
-          pure (baseCtx.withGlobalContext gctx)
+            else gctx.push datatypeName (GlobalKind.type typeParamNames.toList none)
+          -- Add .tvar bindings for type parameters
+          let loc := typeParamsT.info.loc
+          let tctx := typeParamNames.foldl (init := baseCtx.withGlobalContext gctx) fun ctx name =>
+            ctx.push { ident := name, kind := .tvar loc name }
+          pure tctx
         | _, _ => continue
       | none =>
-        match ae.contextLevel with
-        | some idx =>
-          match trees[idx] with
-          | some t => pure t.resultContext
+        -- Check for typeVarsScope (polymorphic function type parameters)
+        match ae.typeVarsScope with
+        | some typeArgsLevel =>
+          match trees[typeArgsLevel] with
+          | some typeArgsTree =>
+            -- Extract identifiers from TypeArgs and add as .tvar bindings
+            let baseCtx := typeArgsTree.resultContext
+            let typeVarNames := extractTypeVarNames typeArgsTree
+            let tctx := typeVarNames.foldl (init := baseCtx) fun ctx name =>
+              let loc := typeArgsTree.info.loc
+              ctx.push { ident := name, kind := .tvar loc name }
+            pure tctx
           | none => continue
-        | none => pure tctx0
+        | none =>
+          match ae.contextLevel with
+          | some idx =>
+            match trees[idx] with
+            | some t => pure t.resultContext
+            | none => continue
+          | none => pure tctx0
     let astx := args[ae.syntaxLevel]
     let expectedKind := argDecls[argLevel].kind
     match expectedKind with
@@ -1209,6 +1281,9 @@ partial def elabExpr (tctx : TypingContext) (stx : Syntax) : ElabM Tree := do
       | .type _ _params _ =>
         logErrorMF loc mf!"{name} is a type when an expression is required."
         return default
+      | .tvar _ _ =>
+        logErrorMF loc mf!"{name} is a type variable when an expression is required."
+        return default
       | .cat c =>
         logErrorMF loc mf!"{name} has category {c} when an expression is required."
         return default

Strata/DDM/Elab/DialectM.lean

@@ -29,6 +29,7 @@ private def foldBoundTypeVars {α} (tp : PreType) (init : α) (f : α → Nat 
   | .ident _ _ a => a.attach.foldl (init := init) fun r ⟨e, _⟩ => e.foldBoundTypeVars r f
   | .fvar _ _ a => a.attach.foldl (init := init) fun r ⟨e, _⟩ => e.foldBoundTypeVars r f
   | .bvar _ i => f init i
+  | .tvar _ _ => init
   | .arrow _ a r => r.foldBoundTypeVars (a.foldBoundTypeVars init f) f
   | .funMacro _ _ r => r.foldBoundTypeVars init f
 

Strata/DDM/Elab/SyntaxElab.lean

@@ -27,6 +27,9 @@ structure ArgElaborator where
   -- Datatype scope: (nameLevel, typeParamsLevel) for recursive datatype definitions
   -- When set, the datatype name is added to the typing context as a type
   datatypeScope : Option (Fin argLevel × Fin argLevel) := .none
+  -- Type vars scope: level of TypeArgs argument for polymorphic function declarations
+  -- When set, identifiers from TypeArgs are added as .tvar bindings
+  typeVarsScope : Option (Fin argLevel) := .none
   -- Whether to unwrap this argument
   unwrap : Bool := false
 deriving Inhabited, Repr
@@ -65,6 +68,7 @@ private def push (as : ArgElaborators)
     argLevel := argLevel.val
     contextLevel := argDecls.argScopeLevel argLevel
     datatypeScope := argDecls.argScopeDatatypeLevel argLevel
+    typeVarsScope := argDecls.argScopeTypeVarsLevel argLevel
   }
   have scp : sc < sc + 1 := by grind
   { as with argElaborators := as.argElaborators.push ⟨newElab, scp⟩ }
@@ -81,6 +85,7 @@ private def pushWithUnwrap
     argLevel := argLevel.val
     contextLevel := argDecls.argScopeLevel argLevel
     datatypeScope := argDecls.argScopeDatatypeLevel argLevel
+    typeVarsScope := argDecls.argScopeTypeVarsLevel argLevel
     unwrap := unwrap
   }
   have scp : sc < sc + 1 := by grind

Strata/DDM/Elab/Tree.lean

@@ -32,6 +32,11 @@ over the parameters.
 Variable belongs to the particular category below.
 -/
 | cat (k : SyntaxCat)
+/--
+Variable is a polymorphic type variable (for function type parameters).
+These are passed through to the dialect's typechecker for inference.
+-/
+| tvar (ann : SourceRange) (name : String)
 deriving Inhabited, Repr
 
 namespace BindingKind
@@ -48,13 +53,15 @@ def ofCat (c : SyntaxCat) : BindingKind :=
 def categoryOf : BindingKind → SyntaxCat
 | .expr tp => .atom tp.ann q`Init.Expr
 | .type loc _ _ => .atom loc q`Init.Type
+| .tvar loc _ => .atom loc q`Init.Type
 | .cat c => c
 
 instance : ToStrataFormat BindingKind where
   mformat bk := private
     match bk with
     | .expr tp => mformat tp
     | .type _ params _ => mformat (params.foldr (init := f!"Type") (fun a f => f!"({a} : Type) -> {f}"))
+    | .tvar _ name => mformat f!"tvar({name})"
     | .cat c => mformat c
 
 end BindingKind