Size-aware base-case weights + RecAwareGradedScore

Specimen/DeriveConstrainedProducer.lean

@@ -896,21 +896,25 @@ def deriveConstrainedProducer
     constrainingInductive inductiveLevels freshArgIdents freshenedOutputNames.toList
     outputTypes.toList producerSort localCtx
 
-/-- Compile a schedule to a weighted sub-producer term. Handles the common pattern of:
-    schedule → MExp → TSyntax, then wrapping with the weight function for the backtracking
-    combinator. Returns the compiled term and whether it should go in the recursive bucket. -/
-private def compileWeightedProducer
+/-- Compile a schedule to a sub-producer term (no weight wrapper).
+    Returns the compiled generator/enumerator/checker body. -/
+private def compileSubProducer
     (schedule : List ScheduleStep × ScheduleSort)
     (outputType : Expr) (deriveSort : DeriveSort)
-    (fuelPrimeName sizePrimeName targetInductive : Name)
-    (weightFnIdent : Ident) (modifierIdent : Option Ident)
-    (ctorName : Name) (outputIndices : List Nat)
-    (badness : Float) (isRecursive : Bool)
-    (freshSize' numBaseLit numRecLit : TSyntax `term) : TermElabM (TSyntax `term) := do
+    (fuelPrimeName sizePrimeName targetInductive : Name) : TermElabM (TSyntax `term) := do
   let (subProducer, _) ← StateT.run (s := #[]) (do
     let mexp ← MExp.scheduleToMExp schedule (.MId `size) (.MId `initSize) outputType
       (fuelPrimeName := fuelPrimeName) (sizePrimeName := sizePrimeName) (targetInductive := targetInductive)
     MExp.mexpToTSyntax mexp deriveSort)
+  return subProducer
+
+/-- Wrap a compiled sub-producer with a weight annotation for the backtracking combinator. -/
+private def wrapWithWeight
+    (subProducer : TSyntax `term) (deriveSort : DeriveSort)
+    (weightFnIdent : Ident) (modifierIdent : Option Ident)
+    (ctorName : Name) (outputIndices : List Nat)
+    (badness : Float) (isRecursive : Bool)
+    (sizeTerm numBaseLit numRecLit : TSyntax `term) : TermElabM (TSyntax `term) := do
   let badnessLit := Syntax.mkScientificLit (toString badness)
   let ctorNameLit := Lean.quote ctorName
   let outputIndicesLit := Lean.quote outputIndices
@@ -921,22 +925,30 @@ private def compileWeightedProducer
     | .Theorem => `(Schedules.DeriveSort.Theorem)
   match deriveSort with
   | .Generator =>
-    let baseWeight ←
-      if isRecursive then
-        `($weightFnIdent $ctorNameLit $outputIndicesLit $deriveSortLit $badnessLit true $freshSize' $numBaseLit $numRecLit)
-      else
-        `($weightFnIdent $ctorNameLit $outputIndicesLit $deriveSortLit $badnessLit false 0 $numBaseLit $numRecLit)
+    let isRecLit := Lean.quote isRecursive
+    let baseWeight ← `($weightFnIdent $ctorNameLit $outputIndicesLit $deriveSortLit $badnessLit $isRecLit $sizeTerm $numBaseLit $numRecLit)
     let finalWeight ← match modifierIdent with
       | none => pure baseWeight
       | some modIdent =>
-        if isRecursive then
-          `($modIdent $baseWeight $ctorNameLit $outputIndicesLit $deriveSortLit $badnessLit true $freshSize' $numBaseLit $numRecLit)
-        else
-          `($modIdent $baseWeight $ctorNameLit $outputIndicesLit $deriveSortLit $badnessLit false 0 $numBaseLit $numRecLit)
+        `($modIdent $baseWeight $ctorNameLit $outputIndicesLit $deriveSortLit $badnessLit $isRecLit $sizeTerm $numBaseLit $numRecLit)
     `( ($finalWeight, $subProducer) )
   | .Enumerator => pure subProducer
   | .Checker | .Theorem => `(fun (_ : Unit) => $subProducer)
 
+/-- Compile a schedule to a weighted sub-producer term. Handles the common pattern of:
+    schedule → MExp → TSyntax, then wrapping with the weight function for the backtracking
+    combinator. Returns the compiled term and whether it should go in the recursive bucket. -/
+private def compileWeightedProducer
+    (schedule : List ScheduleStep × ScheduleSort)
+    (outputType : Expr) (deriveSort : DeriveSort)
+    (fuelPrimeName sizePrimeName targetInductive : Name)
+    (weightFnIdent : Ident) (modifierIdent : Option Ident)
+    (ctorName : Name) (outputIndices : List Nat)
+    (badness : Float) (isRecursive : Bool)
+    (freshSize' numBaseLit numRecLit : TSyntax `term) : TermElabM (TSyntax `term) := do
+  let subProducer ← compileSubProducer schedule outputType deriveSort fuelPrimeName sizePrimeName targetInductive
+  wrapWithWeight subProducer deriveSort weightFnIdent modifierIdent ctorName outputIndices badness isRecursive freshSize' numBaseLit numRecLit
+
 /-- Walk a ConstructorExpr tree and collect which type params it references. -/
 def extractTypeParamRefs (typeParams : Std.HashSet Name) : ConstructorExpr → Std.HashSet Name
   | .Unknown n => if typeParams.contains n then Std.HashSet.ofList [n] else {}
@@ -1283,15 +1295,30 @@ def compileInductiveSchedule (indSched : InductiveSchedule)
     let numRec := indSched.recSchedules.length + numBaseMutual.length
     let numBaseLit := Syntax.mkNumLit (toString numBase)
     let numRecLit := Syntax.mkNumLit (toString numRec)
+    -- Non-recursive producers get different weight terms in the two size branches:
+    -- - size=0 branch (baseProducers): size' is unbound, use literal 0
+    -- - size>0 branch (inductiveProducers): size' is bound, pass it through
+    -- The sub-producer body is compiled once and wrapped with weights for each branch.
+    let mut nonRecProducersForBase : Array (TSyntax `term) := #[]
+    let freshSizeZero ← `((0 : Nat))
     for (ctorName, schedule) in indSched.baseSchedules do
       let (steps, sort) := schedule
       let rewrittenSteps := rewriteSchedule steps
       let isRec := scheduleUsesMutualCall rewrittenSteps
-      let term ← compileWeightedProducer (rewrittenSteps, sort) outputType key.deriveSort
-        freshFuelPrimeName freshSizePrimeName key.inductiveName
-        weightFnIdent modifierIdent ctorName key.outputIndices (lookupCtorBadness ctorName) isRec freshSize' numBaseLit numRecLit
-      if isRec then recursiveProducers := recursiveProducers.push term
-      else nonRecursiveProducers := nonRecursiveProducers.push term
+      if isRec then
+        let term ← compileWeightedProducer (rewrittenSteps, sort) outputType key.deriveSort
+          freshFuelPrimeName freshSizePrimeName key.inductiveName
+          weightFnIdent modifierIdent ctorName key.outputIndices (lookupCtorBadness ctorName) true freshSize' numBaseLit numRecLit
+        recursiveProducers := recursiveProducers.push term
+      else
+        let subProducer ← compileSubProducer (rewrittenSteps, sort) outputType key.deriveSort
+          freshFuelPrimeName freshSizePrimeName key.inductiveName
+        let termInductive ← wrapWithWeight subProducer key.deriveSort
+          weightFnIdent modifierIdent ctorName key.outputIndices (lookupCtorBadness ctorName) false freshSize' numBaseLit numRecLit
+        nonRecursiveProducers := nonRecursiveProducers.push termInductive
+        let termBase ← wrapWithWeight subProducer key.deriveSort
+          weightFnIdent modifierIdent ctorName key.outputIndices (lookupCtorBadness ctorName) false freshSizeZero numBaseLit numRecLit
+        nonRecProducersForBase := nonRecProducersForBase.push termBase
     for (ctorName, schedule) in indSched.recSchedules do
       let (steps, sort) := schedule
       let term ← compileWeightedProducer (rewriteSchedule steps, sort) outputType key.deriveSort
@@ -1306,13 +1333,13 @@ def compileInductiveSchedule (indSched : InductiveSchedule)
         match key.deriveSort with
         | .Checker | .Theorem =>
           let failsafe ← `((fun (_ : Unit) => $failFn $genericFailure))
-          pure (nonRecursiveProducers.push failsafe)
+          pure (nonRecProducersForBase.push failsafe)
         | .Enumerator =>
           let failsafe ← `($failFn $genericFailure)
-          pure (nonRecursiveProducers.push failsafe)
-        | .Generator => pure nonRecursiveProducers
+          pure (nonRecProducersForBase.push failsafe)
+        | .Generator => pure nonRecProducersForBase
       else
-        pure nonRecursiveProducers
+        pure nonRecProducersForBase
     let baseProducers ← `([$baseProducersWithFailsafe,*])
     let allProducers := nonRecursiveProducers ++ recursiveProducers
     let inductiveProducers ← `([$allProducers,*])

Specimen/Scoring.lean

@@ -1454,4 +1454,157 @@ initialize do
   registerScoringBundle { base with wholeScheduleScorer := some sourceQualityWholeScheduleScorer }
 
 
+----------------------------------------------
+-- Built-in: RecAwareGradedScore
+-- Like GradedUniformDensityScore but with a RecursionKind axis that
+-- distinguishes direct recursion (Source.Rec) from same-inductive
+-- cross-mode calls (Source.NonRec targeting the same inductive with
+-- different output indices or derive sort). Direct recursion is
+-- preferred as the simpler pattern when both are available.
+----------------------------------------------
+
+inductive RecursionKind
+  | None
+  | Direct
+  | Mutual
+  deriving Repr, BEq, Inhabited
+
+namespace RecursionKind
+
+def toNat : RecursionKind → Nat
+  | .None => 0
+  | .Direct => 1
+  | .Mutual => 2
+
+instance : Ord RecursionKind where
+  compare a b := compare a.toNat b.toNat
+
+def max (a b : RecursionKind) : RecursionKind :=
+  if a.toNat ≥ b.toNat then a else b
+
+end RecursionKind
+
+structure RecAwareGradedScore where
+  density        : Density := .Total
+  recursionKind  : RecursionKind := .None
+  checkSpeed     : CheckSpeed := .NotACheck
+  passLikelihood : PassLikelihood := .Certain
+  varDeps        : Nat := 0
+  deriving Repr, BEq, Inhabited
+
+deriving instance TypeName for RecAwareGradedScore
+
+instance : Ord RecAwareGradedScore where
+  compare a b :=
+    match compare a.density.toNat b.density.toNat with
+    | .eq => match compare a.recursionKind.toNat b.recursionKind.toNat with
+      | .eq => match compare a.checkSpeed.toNat b.checkSpeed.toNat with
+        | .eq => match compare a.passLikelihood.toNat b.passLikelihood.toNat with
+          | .eq => compare a.varDeps b.varDeps
+          | r => r
+        | r => r
+      | r => r
+    | r => r
+
+instance : LT RecAwareGradedScore := ltOfOrd
+
+instance : Scorable RecAwareGradedScore where
+  empty := {}
+  combine a b :=
+    { density := Density.max a.density b.density
+      recursionKind := RecursionKind.max a.recursionKind b.recursionKind
+      checkSpeed := CheckSpeed.max a.checkSpeed b.checkSpeed
+      passLikelihood := PassLikelihood.max a.passLikelihood b.passLikelihood
+      varDeps := a.varDeps + b.varDeps }
+  isBetter a b := a < b
+  bestOf scores := scores.foldl (fun acc s => if s < acc then s else acc) (scores.headD {})
+  uncoveredPenalty := { density := .Partial, varDeps := 0 }
+  worst := { density := .Checking, recursionKind := .Mutual, checkSpeed := .Recursive, passLikelihood := .Desperate, varDeps := 1000 }
+  badness s :=
+    let varDepPenalty := min 0.05 (s.varDeps.toFloat * 0.01)
+    let recPenalty := s.recursionKind.toNat.toFloat * 0.15
+    if s.density != .Checking then
+      let level := s.density.toNat.toFloat / 4.0
+      min 1.0 (level + recPenalty + varDepPenalty)
+    else
+      let speedVal := match s.checkSpeed with
+        | .NotACheck => 0.0 | .Decidable => 0.0 | .Moderate => 0.33
+        | .Expensive => 0.66 | .Recursive => 1.0
+      let likelVal := match s.passLikelihood with
+        | .Certain => 0.0 | .Likely => 0.0 | .Moderate => 0.33
+        | .Unlikely => 0.66 | .Desperate => 1.0
+      let severity := max speedVal likelVal * 0.7 + min speedVal likelVal * 0.3
+      min 1.0 (0.75 + severity * 0.25 + recPenalty + varDepPenalty)
+
+private def classifyRecursionKind (key : SpecKey)
+    (src : Source) (outputs : List (Name × Option ConstructorExpr)) (prodSort : ProducerSort) : RecursionKind :=
+  match src with
+  | .Rec .. => .Direct
+  | .MutRec .. => .Mutual
+  | .NonRec (indName, args) =>
+    if indName != key.inductiveName then .None
+    else
+      let outputIdxs := outputs.filterMap fun (n, _) =>
+        args.findIdx? fun a => match a with | .Unknown v => v == n | _ => false
+      let depDeriveSort := match prodSort with
+        | .Enumerator => DeriveSort.Enumerator
+        | .Generator => DeriveSort.Generator
+      let depKey : SpecKey := { inductiveName := indName, outputIndices := outputIdxs, deriveSort := depDeriveSort }
+      if depKey == key then .Direct
+      else .Mutual
+
+private def classifyRecursionKindCheck (key : SpecKey) (src : Source) : RecursionKind :=
+  match src with
+  | .Rec .. => .Direct
+  | .MutRec .. => .Mutual
+  | .NonRec (indName, _args) =>
+    if indName != key.inductiveName then .None
+    else
+      if key.deriveSort == .Checker && key.outputIndices.isEmpty then .Direct
+      else .Mutual
+
+def recAwareStepScorer : StepScorer RecAwareGradedScore := fun key memo inputVars step => do
+  match step with
+  | .Unconstrained .. => return { density := .Total }
+  | .Match .. => return { density := .Backtracking }
+  | .Check src polarity =>
+    let varDeps := countGeneratedVarDeps inputVars src
+    let speed := classifyCheckSpeed memo key src varDeps
+    let likelihood ← classifyPassLikelihood inputVars key src polarity varDeps
+    let recKind := classifyRecursionKindCheck key src
+    return { density := .Checking, recursionKind := recKind, checkSpeed := speed, passLikelihood := likelihood, varDeps := varDeps }
+  | .SuchThat outputs src prodSort =>
+    let outputNames := Std.HashSet.ofList (outputs.map (·.1))
+    let varDeps := countGeneratedVarDeps (inputVars.union outputNames) src
+    let recKind := classifyRecursionKind key src outputs prodSort
+    let depDensity : Density := match src with
+      | .Rec .. => .Partial
+      | .MutRec .. => .Partial
+      | .NonRec (indName, args) =>
+        let outputIdxs := outputs.filterMap fun (n, _) =>
+          args.findIdx? fun a => match a with | .Unknown v => v == n | _ => false
+        let depDeriveSort := match prodSort with
+          | .Enumerator => DeriveSort.Enumerator
+          | .Generator => DeriveSort.Generator
+        let depKey : SpecKey := { inductiveName := indName, outputIndices := outputIdxs, deriveSort := depDeriveSort }
+        if depKey == key then .Partial
+        else match memo[depKey]? with
+          | some (.done depSched) => (Score.unwrap RecAwareGradedScore depSched.score).density
+          | _ => .Partial
+    return { density := depDensity, recursionKind := recKind, varDeps := varDeps }
+
+def recAwareScheduleScorer : ScheduleScorer RecAwareGradedScore := fun stepScores =>
+  stepScores.foldl Scorable.combine Scorable.empty
+
+def recAwareLeafAggregator : LeafAggregator RecAwareGradedScore := fun ctors =>
+  match ctors with
+  | [] => Scorable.uncoveredPenalty
+  | _ => Scorable.bestOf (ctors.map Prod.snd)
+
+def recAwareInductiveAggregator : InductiveAggregator RecAwareGradedScore := fun leafScores =>
+  leafScores.foldl Scorable.combine Scorable.empty
+
+initialize registerScoringBundle (mkScorerBundle `Scoring.RecAwareGradedScore
+  recAwareStepScorer recAwareScheduleScorer recAwareLeafAggregator recAwareInductiveAggregator)
+
 end Scoring