Support structure-parameterized output types in the constrained deriver

Specimen/DeriveConstrainedProducer.lean

@@ -265,10 +265,14 @@ def getScheduleSort (conclusion : HypothesisExpr)
     - Note: it is the caller's responsibility to check that `conclusion` does indeed contain
       a non-trivial function application (e.g. by using `containsNonTrivialFuncApp`) -/
 def linearizeAndFlatten
-  (hypotheses : Array Expr) (conclusion : Expr) (outputIndices : List Nat) (localCtx : LocalContext) :
+  (hypotheses : Array Expr) (conclusion : Expr) (outputIndices : List Nat) (localCtx : LocalContext)
+  (fixedFVars : Std.HashSet FVarId := {}) :
   UnifyM (Array Expr × Expr × List (Name × Expr) × LocalContext) := do
-  -- Phase 1: flatten function calls into fresh unknowns with equality hypotheses
-  let funcAppExprs ← collectUnmatchableProperSubterms conclusion
+  -- Phase 1: flatten function calls into fresh unknowns with equality hypotheses.
+  -- `fixedFVars` are the producer's fixed inputs; a subterm determined entirely
+  -- by them (e.g. a type parameter `T.mono` where `T` is an input) is left in
+  -- place rather than lifted into a generated unknown.
+  let funcAppExprs ← collectUnmatchableProperSubterms fixedFVars conclusion
   trace[plausible.deriving.arbitrary] m!"Unmatchable exprs: {funcAppExprs} In conclusion: {conclusion}"
   withLCtx' localCtx do
 
@@ -473,8 +477,22 @@ def getScheduleForInductiveRelationConstructor
     -- equal to the result of the function call, and adding an extra hypothesis asserting equality
     -- between the function call and the variable.
     -- `freshNamesAndTypes` is a list containing the names & types of the fresh variables produced during this procedure.
+    -- Identify fixed inputs as fvars, so flattening can leave subterms that are
+    -- determined entirely by the inputs (e.g. an output type's structure
+    -- parameter `T.mono`) in place instead of trying to generate them. The
+    -- conclusion's arguments at non-output (input) positions that are bare
+    -- variables are exactly those inputs (e.g. the inductive's parameter `T`).
+    let curLCtx ← getLCtx
+    let conclArgs := conclusion.getAppArgs
+    let fixedFVars : Std.HashSet FVarId := Id.run do
+      let mut s : Std.HashSet FVarId := {}
+      for h : i in [:conclArgs.size] do
+        if i ∉ outputIndices then
+          if let .fvar fid := conclArgs[i] then
+            s := s.insert fid
+      return s
     let (updatedHypotheses, updatedConclusion, freshNamesAndTypes, updatedLocalCtx) ←
-      linearizeAndFlatten hypotheses conclusion outputIndices (← getLCtx)
+      linearizeAndFlatten hypotheses conclusion outputIndices curLCtx fixedFVars
     -- Enter the updated `LocalContext` containing the fresh variable that was created when rewriting the conclusion
     withLCtx' updatedLocalCtx (do
       let hypothesisExprs := (← monadLift (updatedHypotheses.toList.mapM (exprToHypothesisExpr ctorName))).toArray

Specimen/Enumerators.lean

@@ -128,6 +128,10 @@ end EnumeratorCombinators
 
 -- Some simple `Enum` instances
 
+/-- `Enum` instance for `Unit` (its single value) -/
+instance : Enum Unit where
+  enum := pureEnum ()
+
 /-- `Enum` instance for `Bool` -/
 instance : Enum Bool where
   enum := pureEnum false <|> pureEnum true

Specimen/MExp.lean

@@ -177,6 +177,38 @@ partial def constructorExprToMExp (exp : Explicit) (expr : ConstructorExpr) : ME
   | .Lit l => .MLit l
   | .CSort lvl => .MSort lvl
 
+/-- Recursively drop the arguments of every *data constructor* application in a
+    `ConstructorExpr` that sit at implicit / instance-implicit positions.
+
+    A conclusion output is emitted in implicit-allowing form, so a constructor's
+    implicit arguments — most importantly an output inductive's structure
+    type-parameter, e.g. `LExpr.const`'s `T` in `LExpr.const (T.mono) m c`, but
+    also ordinary ones like `Option.some`'s `α` — must be omitted and left for
+    Lean to infer from the explicit arguments. Constructors whose argument count
+    does not match their arity are left unchanged. `FuncApp`/`TyCtor` nodes are
+    traversed but their own argument lists are not filtered (function/type-former
+    applications are emitted implicit-allowing already). -/
+partial def dropImplicitCtorArgsExpr (ce : ConstructorExpr) : MetaM ConstructorExpr := do
+  match ce with
+  | .Ctor c args =>
+    let args ← args.mapM dropImplicitCtorArgsExpr
+    -- Only genuine data constructors have implicit-position args to drop; some
+    -- `.Ctor` nodes are actually abbreviations/defs (e.g. `LExprParams.mono`).
+    unless (← getEnv).isConstructor c do return .Ctor c args
+    let ctorType := (← getConstInfoCtor c).type
+    let argsArr := args.toArray
+    let kept ← Meta.forallTelescopeReducing ctorType fun bvars _ => do
+      if bvars.size ≠ argsArr.size then return args
+      let mut result := #[]
+      for h : i in [:argsArr.size] do
+        if (← bvars[i]!.fvarId!.getDecl).binderInfo.isExplicit then
+          result := result.push argsArr[i]!
+      return result.toList
+    return .Ctor c kept
+  | .TyCtor c args => return .TyCtor c (← args.mapM dropImplicitCtorArgsExpr)
+  | .FuncApp f args => return .FuncApp f (← args.mapM dropImplicitCtorArgsExpr)
+  | other => return other
+
 partial def mexpToConstructorExpr (m : MExp) : Option ConstructorExpr :=
   match m with
   | .MId u => return .Unknown u
@@ -475,7 +507,7 @@ def scheduleStepToMExp (step : ScheduleStep) (defFuel : MExp) (k : MExp) (output
       for the function we are deriving (these correspond to `size` and `initSize`
       in the QuickChick code for the derived functions)
     - `targetInductive` is the inductive we're generating for (used to detect same-type deps) -/
-def scheduleToMExp (schedule : Schedule) (mfuel : MExp) (defFuel : MExp) (recType : Expr) (fuelPrimeName : Name := `fuel') (sizePrimeName : Name := `size') (targetInductive : Name := `_unknown) : CompileScheduleM MExp :=
+def scheduleToMExp (schedule : Schedule) (mfuel : MExp) (defFuel : MExp) (recType : Expr) (fuelPrimeName : Name := `fuel') (sizePrimeName : Name := `size') (targetInductive : Name := `_unknown) : CompileScheduleM MExp := do
   let (scheduleSteps, scheduleSort) := schedule
   -- Compute the size expression: split budget across all size-consuming calls
   -- (self-recursive, mutual, and non-recursive calls to the same inductive)
@@ -485,25 +517,27 @@ def scheduleToMExp (schedule : Schedule) (mfuel : MExp) (defFuel : MExp) (recTyp
     else .MApp .allExplicit (.MConst ``Nat.div) [.MId sizePrimeName, .MLit (.natVal numSizeCalls)]
   -- Determine the *epilogue* of the schedule (i.e. what happens after we
   -- have finished executing all the `scheduleStep`s)
-  let epilogue :=
+  let epilogue ← do
     match scheduleSort with
     | .ProducerSchedule _ conclusionOutputs =>
-      -- Convert all the outputs in the conclusion to `mexp`s
+      -- Drop implicit constructor arguments (e.g. an output type's structure
+      -- parameter), then convert all the outputs in the conclusion to `mexp`s.
+      let conclusionOutputs ← conclusionOutputs.mapM (fun ce => (monadLift (dropImplicitCtorArgsExpr ce) : CompileScheduleM _))
       let conclusionMExps := constructorExprToMExp .allowImplicit <$> conclusionOutputs
       -- If there are multiple outputs, wrap them in a tuple
       match conclusionMExps with
       | [] => panic! "No outputs being returned in producer schedule"
-      | [output] => MExp.MRet output
-      | outputs => MExp.MRet (tupleOfList (fun e1 e2 => .MApp .allowImplicit (.MConst ``Prod.mk) [e1, e2]) outputs outputs[0]?)
-    | .CheckerSchedule => okTrue
+      | [output] => pure (MExp.MRet output)
+      | outputs => pure (MExp.MRet (tupleOfList (fun e1 e2 => .MApp .allowImplicit (.MConst ``Prod.mk) [e1, e2]) outputs outputs[0]?))
+    | .CheckerSchedule => pure okTrue
     | .TheoremSchedule conclusion typeClassUsed =>
       -- Create a pattern-match on the result of hte checker
       -- on the conclusion, returning `.ok true` or `.ok false` accordingly
       let conclusionMExp := hypothesisExprToMExp conclusion
       let scrutinee :=
         if typeClassUsed then decOptChecker conclusionMExp mfuel
         else conclusionMExp
-      matchExceptBool scrutinee okTrue okFalse
+      pure (matchExceptBool scrutinee okTrue okFalse)
   -- Fold over the `scheduleSteps` and convert each of them to a functional `MExp`
   -- Note that the fold composes the `MExp`, and we use `foldr` since
   -- we want the `epilogue` to be the base-case of the fold

Specimen/MakeConstrainedProducerInstance.lean

@@ -53,6 +53,62 @@ def mkTypeClassInstanceBinders (typeParams : Array Name) (typeClasses : Array Na
       `(Lean.Elab.Deriving.instBinderF| [$(mkIdent tc) $(mkIdent param)])
   return TSyntaxArray.mk instances
 
+open TSyntax.Compat in
+/-- Recursively expand a *structure-typed* parameter into `[className proj]`
+    instance binders for each leaf field of type `Type u`. Mirrors
+    `expandStructBinders` in `Specimen.DeriveArbitrary` (the unconstrained
+    `deriving Arbitrary` path), which is what lets `Arbitrary (LExpr T)` be
+    derived for a structure parameter `T`: the constrained producer instance
+    needs the same per-field binders (e.g. `[Arbitrary T.base.Metadata]`) so the
+    metadata fields carried by each constructor can be generated.
+
+    `ty` is the type being walked (the parameter's type, or a field's type as we
+    recurse) and `syn` is the matching surface syntax (the parameter identifier,
+    or a projection chain into it). A leaf `Type u` field yields `[className syn]`;
+    a structure field recurses into its own fields. Returns `#[]` for fields that
+    are neither Types nor structures-of-Types — the conservative,
+    behavior-preserving choice for the constrained path. -/
+partial def expandStructInstBinders (className : Name) (ty : Expr) (syn : TSyntax `term) :
+    TermElabM (TSyntaxArray `Lean.Parser.Term.bracketedBinder) := do
+  if ty.isSort then
+    -- A `Type u` leaf: emit `[className syn]`.
+    return #[← `(Lean.Elab.Deriving.instBinderF| [$(mkCIdent className):ident $syn])]
+  let env ← getEnv
+  let some sName := ty.constName? | return #[]
+  let some sInfo := getStructureInfo? env sName | return #[]
+  let mut result : TSyntaxArray `Lean.Parser.Term.bracketedBinder := #[]
+  for field in sInfo.fieldNames do
+    let projName := sName ++ field
+    -- The field's declared type, read as the codomain of the projection's
+    -- signature `∀ (_ : sName ..), fieldType`. Structure fields here are the
+    -- metadata-configuration types, whose field types do not depend on the value.
+    -- Limitation: `mkConst projName` is built without universe-level arguments, so
+    -- for a *universe-polymorphic* structure parameter `projType` would be computed
+    -- at the wrong universe. We don't support such parameters (they are rare in
+    -- practice; all current Strata params are `Type 0`); supporting them would mean
+    -- extracting `ty`'s universe levels and passing them to `mkConst` here.
+    let projType ← forallTelescopeReducing (← inferType (mkConst projName))
+      (fun _ body => pure body)
+    let projSyn ← `($(mkIdent projName) $syn)
+    result := result ++ (← expandStructInstBinders className projType projSyn)
+  return result
+
+/-- Build the unconstrained-producer instance binders for a list of
+    non-target inductive parameters. For a `Sort`-typed parameter `α`, emits
+    `[className α] [DecidableEq α]` (the existing behavior). For a *structure*
+    parameter `T`, expands it into per-field binders via `expandStructInstBinders`
+    (e.g. `[className T.base.Metadata]` …) so the structure's metadata fields can
+    be generated. Parameters of other types contribute no binders. -/
+def mkProducerParamInstBinders (className : Name)
+    (params : Array (Name × Expr)) : TermElabM (TSyntaxArray `Lean.Parser.Term.bracketedBinder) := do
+  let mut result : TSyntaxArray `Lean.Parser.Term.bracketedBinder := #[]
+  for (paramName, paramType) in params do
+    if paramType.isSort then
+      result := result ++ (← mkTypeClassInstanceBinders #[paramName] #[className, ``DecidableEq])
+    else
+      result := result ++ (← expandStructInstBinders className paramType (mkIdent paramName))
+  return result
+
 /-- Finds the index of the argument in the inductive application for the value we wish to generate
     (i.e. finds `i` s.t. `args[i] == targetVar`) -/
 def findTargetVarIndex (targetVar : FVarId) (args : Array Expr) : (Option Nat) := do
@@ -135,6 +191,10 @@ def mkConstrainedProducerTypeClassInstance
     let mut outerParams := #[]
     let mut outputTypeSyntaxes : Array (TSyntax `term) := #[]
     let mut typeParams := #[]
+    -- Non-target, non-sort *structure* parameters (e.g. `T : LExprParams`): we
+    -- need per-field unconstrained-producer instances for these (see
+    -- `mkProducerParamInstBinders`).
+    let mut structParams : Array (Name × Expr) := #[]
     for (paramName, paramType, paramTypeSyntax) in paramInfo do
       -- Only add a function parameter if the argument to the inductive relation is not a target variable
       -- (We skip the target variables since those are the values we wish to generate)
@@ -152,6 +212,8 @@ def mkConstrainedProducerTypeClassInstance
           `(Term.letIdBinder| ($innerParamIdent : $paramTypeSyntax))
 
         innerParams := innerParams.push innerParam
+        if !paramType.isSort then
+          structParams := structParams.push (paramName, paramType)
       else
         outputTypeSyntaxes := outputTypeSyntaxes.push paramTypeSyntax
 
@@ -203,7 +265,11 @@ def mkConstrainedProducerTypeClassInstance
       | .Generator => ``Plausible.Arbitrary
       | .Enumerator => ``Enum
 
-    let arbitraryTypeParamInstances ← mkTypeClassInstanceBinders typeParams #[producerUnconstrainedClass, ``DecidableEq]
+    let arbitraryTypeParamInstances0 ← mkTypeClassInstanceBinders typeParams #[producerUnconstrainedClass, ``DecidableEq]
+    -- Per-field unconstrained-producer instances for structure parameters (e.g.
+    -- `[Arbitrary T.base.Metadata]`), so their fields can be generated.
+    let structParamInstances ← mkProducerParamInstBinders producerUnconstrainedClass structParams
+    let arbitraryTypeParamInstances := arbitraryTypeParamInstances0 ++ structParamInstances
 
     let fuelVal := Lean.Option.get (← getOptions) specimen.fuel
     let fuelLit := Syntax.mkNumLit (toString fuelVal)
@@ -264,6 +330,8 @@ def mkConstrainedProducerMutualPieces
     innerParamBinders := innerParamBinders.push (← `(($initSizeIdent : $natIdent)))
     innerParamBinders := innerParamBinders.push (← `(($sizeIdent : $natIdent)))
 
+    -- Non-target, non-sort structure parameters needing per-field producer instances.
+    let mut structParams : Array (Name × Expr) := #[]
     for (paramName, paramType, paramTypeSyntax) in paramInfo do
       if paramType.isSort then
         typeParams := typeParams.push paramName
@@ -274,6 +342,7 @@ def mkConstrainedProducerMutualPieces
           innerParamBinders := innerParamBinders.push (← `(($(mkIdent paramName) : Sort _)))
         else
           innerParamBinders := innerParamBinders.push (← `(($(mkIdent paramName) : $paramTypeSyntax)))
+          structParams := structParams.push (paramName, paramType)
       else
         outputTypeSyntaxes := outputTypeSyntaxes.push paramTypeSyntax
 
@@ -325,11 +394,23 @@ def mkConstrainedProducerMutualPieces
           | .Enumerator => ``Enum
         #[producerUnconstrainedClass, ``DecidableEq]
     let defTypeParamInstances ← mkTypeClassInstanceBinders typeParams typeClasses
+    -- Per-field producer instances for structure parameters (e.g.
+    -- `[Arbitrary T.base.Metadata]`). These reference value params (`T`), so they
+    -- are placed *innermost* — after all value params — where those are in scope.
+    -- The `def` itself always builds the underlying producer, so it uses the
+    -- producer-based unconstrained class regardless of `deriveSort`.
+    let structParamClass := match producerSort with
+      | .Generator => ``Plausible.Arbitrary
+      | .Enumerator => ``Enum
+    let structParamInstances ← mkProducerParamInstBinders structParamClass structParams
 
     -- Emit the def with ∀ type (supports instance binders inline)
     let defIdent := mkIdent globalDefName
     -- Build ∀ type with instance binders interleaved after Sort-typed params
     let mut defType ← pure optionTProducerType
+    -- Innermost: structure-parameter field instances (all value params in scope).
+    for instBinder in structParamInstances.reverse do
+      defType ← `(∀ $instBinder:bracketedBinder, $defType)
     -- Add non-Sort value params (from right)
     let insertIdx := 3 + typeParams.size
     for (name, ty) in allParamNamesAndTypes[insertIdx:].toArray.reverse do
@@ -340,9 +421,12 @@ def mkConstrainedProducerMutualPieces
     -- Add Sort-typed params + fuel/initSize/size (from right)
     for (name, ty) in allParamNamesAndTypes[:insertIdx].toArray.reverse do
       defType ← `(($name : $ty) → $defType)
-    -- Lambda includes instance binders at the same position
+    -- Lambda includes instance binders at the same positions as the ∀ type:
+    -- sort-param instances after the sort params, struct-field instances innermost.
     let instParams : Array (TSyntax `term) := defTypeParamInstances.map (fun b => ⟨b.raw⟩)
-    let allInnerParams := innerParamBinders[:insertIdx].toArray ++ instParams ++ innerParamBinders[insertIdx:].toArray
+    let structInstParams : Array (TSyntax `term) := structParamInstances.map (fun b => ⟨b.raw⟩)
+    let allInnerParams := innerParamBinders[:insertIdx].toArray ++ instParams
+      ++ innerParamBinders[insertIdx:].toArray ++ structInstParams
     let lambdaBody ← `(fun $allInnerParams* => $matchExpr)
     let defCmd ← `(command| def $defIdent : $defType := $lambdaBody)
 
@@ -360,7 +444,10 @@ def mkConstrainedProducerMutualPieces
           | .Enumerator => #[``Enum, ``DecidableEq]
     let instCmd ← match deriveSort with
       | .Checker | .Theorem => do
-        let arbitraryTypeParamInstances ← mkTypeClassInstanceBinders typeParams instTypeClasses
+        let arbitraryTypeParamInstances0 ← mkTypeClassInstanceBinders typeParams instTypeClasses
+        -- Per-field struct-param instances for the wrapper (Checker/Theorem → `Enum`).
+        let structInsts ← mkProducerParamInstBinders ``Enum structParams
+        let arbitraryTypeParamInstances := arbitraryTypeParamInstances0 ++ structInsts
         `(command|
           instance $arbitraryTypeParamInstances:bracketedBinder* : $decOptTypeclass (@$(mkIdent inductiveName) $args*) where
             $unqualifiedDecOptFn:ident := fun $freshSizeIdent => $callExpr)
@@ -371,7 +458,14 @@ def mkConstrainedProducerMutualPieces
         let producerTypeClassFunction := match producerSort with
           | .Generator => unqualifiedArbitrarySizedSTFn
           | .Enumerator => unqualifiedEnumSizedSTFn
-        let arbitraryTypeParamInstances ← mkTypeClassInstanceBinders typeParams instTypeClasses
+        let producerUnconstrainedClass := match producerSort with
+          | .Generator => ``Plausible.Arbitrary
+          | .Enumerator => ``Enum
+        let arbitraryTypeParamInstances0 ← mkTypeClassInstanceBinders typeParams instTypeClasses
+        -- Per-field struct-param instances for the wrapper (Generator/Enumerator
+        -- → `Arbitrary`/`Enum`).
+        let structInsts ← mkProducerParamInstBinders producerUnconstrainedClass structParams
+        let arbitraryTypeParamInstances := arbitraryTypeParamInstances0 ++ structInsts
         `(command|
           instance $arbitraryTypeParamInstances:bracketedBinder* : $producerTypeClass $targetTypeSyntax (fun $targetVarPattern => @$(mkIdent inductiveName) $args*) where
             $producerTypeClassFunction:ident := fun $freshSizeIdent => $callExpr)