import Specimen.UnificationMonad import Specimen.Utils import Specimen.Score open Lean namespace Schedules ---------------------------------------------- -- Type definitions ---------------------------------------------- /-- A `HypothesisExpr` is a datatype representing a hypothesis for a constructor of an inductive relation, consisting of a constructor name applied to some list of arguments, each of which are `ConstructorExpr`s -/ abbrev HypothesisExpr := Name × List ConstructorExpr /-- `ToExpr` instance for `HypothesisExpr` (for converting `HypothesisExpr`s to Lean `Expr`s) -/ instance : ToExpr HypothesisExpr where toExpr (hypExpr : HypothesisExpr) : Expr := let (ctorName, ctorArgs) := hypExpr mkAppN (mkConst ctorName) (toExpr <$> ctorArgs.toArray) toTypeExpr := mkConst ``Expr local instance [Ord α][Ord β]: Ord (α × β) := lexOrd /-- A source is the thing we wish to check/generate/enumerate -/ inductive Source | NonRec : HypothesisExpr → Source | Rec : Name → List ConstructorExpr → Source | MutRec : Name → List ConstructorExpr → Source deriving Repr, BEq /-- Producers are either enumerators or generators -/ inductive ProducerSort | Enumerator | Generator deriving Repr, BEq, Ord /-- The sort of function we are deriving based on an inductive relation: determines whether we are deriving a (constrained) generator, enumerator or a checker. Note: the `Theorem` constructor is used in the artifact of "Testing Theorems, Fully Automatically" for automatically testing whether a theorem holds (we replicate it here for completeness). -/ inductive DeriveSort | Generator | Enumerator | Checker | Theorem deriving Repr, BEq, Ord, Hashable, Inhabited /-- Determines if a `DeriveSort` corresponds to a producer (only generators & enumerators are considered producers) -/ def DeriveSort.isProducer : DeriveSort → Bool | .Generator | .Enumerator => true | .Checker | .Theorem => false /-- The type of schedule we wish to derive -/ inductive ScheduleSort /-- tuple of produced outputs from conclusion of constructor -/ | ProducerSchedule (producerSort : ProducerSort) (conclusion : List ConstructorExpr) /-- checkers need not bother with conclusion of constructor, only hypotheses need be checked and conclusion of constructor follows-/ | CheckerSchedule /-- In a `TheoremSchedule`, we check the `conclusion` of the theorem, and take in a `Bool` which is true if we need to find a checker by identifying the `DecOpt` instance, and false otherwise (we're currently dealing with a function that returns `Option Bool`) -/ | TheoremSchedule (conclusion : HypothesisExpr) (typeClassUsed : Bool) deriving Repr, BEq inductive Explicit where | allExplicit | allowImplicit deriving Repr, BEq, Ord /-- A single step in a generator schedule -/ inductive ScheduleStep /-- Unconstrained generation -/ | Unconstrained : Name → Source → ProducerSort → ScheduleStep /-- Generate a value such that a predicate is satisfied -/ | SuchThat : List (Name × Option ConstructorExpr) → Source → ProducerSort → ScheduleStep /-- Check whether some proposition holds (the bool is the desired truth value of the proposition we're checking) -/ | Check : Source → Bool → ScheduleStep /-- Used when you decompose a constructor constrained arg into a fresh variable followed by a pattern match -/ | Match : Explicit → Name → Pattern → ScheduleStep deriving Repr, BEq /-- Pretty-print a ConstructorExpr in readable form -/ partial def ppConstructorExpr : ConstructorExpr → String | .Unknown name => name.toString | .Hole => "_" | .Lit (.natVal n) => toString n | .Lit (.strVal s) => s!"\"{s}\"" | .CSort _ => "Sort" | .Ctor ``Prod.mk [_, _, a, b] => s!"({ppConstructorExpr a}, {ppConstructorExpr b})" | .Ctor ``Prod.mk [a, b] => s!"({ppConstructorExpr a}, {ppConstructorExpr b})" | .Ctor ``List.nil _ => "[]" | .Ctor ``List.cons [_, h, t] => s!"{ppConstructorExpr h} :: {ppConstructorExpr t}" | .Ctor ``List.cons [h, t] => s!"{ppConstructorExpr h} :: {ppConstructorExpr t}" | .Ctor ``Bool.true [] => "true" | .Ctor ``Bool.false [] => "false" | .Ctor name args => let shortName := name.componentsRev.head?.getD name |>.toString if args.isEmpty then shortName else s!"{shortName} {" ".intercalate (args.map ppConstructorExpr)}" | .TyCtor name args => let shortName := name.componentsRev.head?.getD name |>.toString if args.isEmpty then shortName else s!"{shortName} {" ".intercalate (args.map ppConstructorExpr)}" | .FuncApp name args => let shortName := name.componentsRev.head?.getD name |>.toString if args.isEmpty then shortName else s!"({shortName} {" ".intercalate (args.map ppConstructorExpr)})" /-- Pretty-print a HypothesisExpr -/ def ppHypothesisExpr (hyp : HypothesisExpr) : String := let (name, args) := hyp if name == ``Eq then match args with | [_, lhs, rhs] => s!"{ppConstructorExpr lhs} = {ppConstructorExpr rhs}" | [lhs, rhs] => s!"{ppConstructorExpr lhs} = {ppConstructorExpr rhs}" | _ => s!"Eq {" ".intercalate (args.map ppConstructorExpr)}" else let shortName := name.componentsRev.head?.getD name |>.toString if args.isEmpty then shortName else s!"{shortName} {" ".intercalate (args.map ppConstructorExpr)}" /-- Stringifier for `Source` -/ def ppSource source := match source with | Source.Rec name ctrArgs => let shortName := name.componentsRev.head?.getD name |>.toString s!"{shortName} {" ".intercalate (ctrArgs.map ppConstructorExpr)}" | Source.MutRec name ctrArgs => let shortName := name.componentsRev.head?.getD name |>.toString s!"⟳ {shortName} {" ".intercalate (ctrArgs.map ppConstructorExpr)}" | Source.NonRec hyp => ppHypothesisExpr hyp def ppPattern pat := ppConstructorExpr (constructorExprOfPattern pat) /-- Stringifier for `step` -/ def ppStep step := match step with | ScheduleStep.Unconstrained name src _ => s!"{name} ← {ppSource src}" | .SuchThat vars src _ => let varNames := vars.map (fun (n : Name × Option ConstructorExpr) => ToString.toString n.1) s!"[{", ".intercalate varNames}] ← {ppSource src}" | .Check src true => s!"check {ppSource src}" | .Check src false => s!"check ¬({ppSource src})" | .Match _ name pattern => s!"match {name} with {ppPattern pattern}" /-- Stringifier for lists of steps (without conclusion). -/ def ppScheduleSteps (steps : List ScheduleStep) := "do\n " ++ String.intercalate "\n " (ppStep <$> steps) /-- Quality score for a generator schedule. Lower is better (ordered lexicographically). A schedule with fewer checks is preferred; among equal check counts, shorter is better; among equal lengths, fewer unconstrained (arbitrary) bindings is better. -/ structure ScheduleScore where checks : Nat length : Nat unconstrained : Nat deriving Ord, Repr def scheduleStepsScore (schedule : List ScheduleStep) : ScheduleScore := Id.run do let mut checks := 0 let mut length := 0 let mut unconstrained := 0 for step in schedule do length := length + 1 match step with | .Check .. => checks := checks + 1 | .Unconstrained .. => unconstrained := unconstrained + 1 | _ => () ⟨checks, length, unconstrained⟩ instance : LT ScheduleScore := ltOfOrd def scheduleLT (a b : List ScheduleStep) := scheduleStepsScore a < scheduleStepsScore b /-- A schedule is a pair consisting of an ordered list of `ScheduleStep`s, and the sort of schedule we're dealing with (the latter is the "conclusion" of the schedule) -/ abbrev Schedule := List ScheduleStep × ScheduleSort /-- Stringifier for a full schedule (steps + conclusion). -/ def ppSchedule (schedule : Schedule) : String := let (steps, sort) := schedule let stepsStr := String.intercalate "\n " (ppStep <$> steps) let conclusionStr := match sort with | .ProducerSchedule _ conclusion => let outputStr := match conclusion with | [e] => ppConstructorExpr e | es => s!"({String.intercalate ", " (es.map ppConstructorExpr)})" s!"\n return {outputStr}" | .CheckerSchedule => s!"\n return ok" | .TheoremSchedule hyp _ => s!"\n check_conclusion {ppHypothesisExpr hyp}" s!"do\n {stepsStr}{conclusionStr}" /-- Each `ScheduleStep` is associated with a `Density`, which represents a failure mode of a generator -/ inductive Density /-- Invokes a call to a checker -/ | Checking /-- A call to `ArbitrarySuchThat`, followed by a pattern-match on the generated value (this happens when we want the output of the generator to have a certain shape) -/ | Backtracking /-- a call to `ArbitrarySuchThat` ??? -/ | Partial /-- Unconstrained generation, i.e. calls to `arbitrary` -/ | Total deriving Repr, BEq /-- Converts a `HypothesisExpr` to a `TSyntax term` -/ def hypothesisExprToTSyntaxTerm (hypExpr : HypothesisExpr) : MetaM (TSyntax `term) := do let (ctorName, ctorArgs) := hypExpr if ctorName = `sort then `(Sort _) else let ctorArgTerms ← ctorArgs.toArray.mapM constructorExprToTSyntaxTerm `($(mkIdent ctorName) $ctorArgTerms:term*) mutual /-- Converts an `Expr` to a `ConstructorExpr` ("classifying" the `Expr`: deciding which `ConstructorExpr` variant — variable, constructor, type constructor, function application, literal, etc. — corresponds to it). -/ partial def exprToConstructorExpr (e : Expr) : MetaM ConstructorExpr := do match e with | .fvar id => let localDecl ← FVarId.getDecl id return ConstructorExpr.Unknown localDecl.userName | .const name _ => -- Check if this is a constructor let env ← getEnv if env.isConstructor name then return ConstructorExpr.Ctor name [] else if (← isInductive name) then return .TyCtor name [] else return .FuncApp name [] | .app .. => do -- Classify the application head, then its arguments. We process arguments -- positionally (rather than via the application spine) so we can consult -- each argument's binder info: an argument in an implicit or -- instance-implicit position that cannot itself be classified (e.g. the -- `GetElem?` validity lambda `fun as i => i < as.length`) is replaced by a -- hole, since it is inferable from the explicit arguments. Arguments in -- explicit positions must still classify, as before. let headExpr ← exprToConstructorExpr e.getAppFn let argExprs ← classifyAppArgs e match headExpr with | .TyCtor name args => return .TyCtor name (args ++ argExprs) | .Ctor name args => return .Ctor name (args ++ argExprs) | .FuncApp name args => return .FuncApp name (args ++ argExprs) | .Unknown name => throwError m!"exprToConstructorExpr: We do not support higher order application of {name} in Expr {e}" | .Lit _ => throwError m!"exprToConstructorExpr: String and Nat Literals cannot be applied as functions, see: {e.getAppFn} in {e}" | .CSort _lvl => throwError m!"exprToConstructorExpr: String and Nat Literals cannot be applied as functions, see: {e.getAppFn} in {e}" | .Hole => -- Unreachable: `exprToConstructorExpr` never yields a hole for an -- application head (holes only arise for implicit/instance *arguments* -- in `classifyAppArgs`). A hole applied as a function is nonsensical. throwError m!"exprToConstructorExpr: a hole cannot be applied as a function, see: {e.getAppFn} in {e}" | .lit l => return .Lit l | .sort lvl => return .CSort lvl | .lam .. => -- Check if this lambda is a typeclass instance (e.g., eta-expanded DecidableEq). -- If so, treat it as an opaque constant — it will be resolved by typeclass -- synthesis in the generated code. Otherwise, reject it. let ty ← Meta.inferType e let headTy := ty.getForallBody if headTy.isApp then let className := headTy.getAppFn.constName! if Lean.isClass (← getEnv) className then return .FuncApp (← Core.mkFreshUserName `inst) [] throwError m!"exprToConstructorExpr can only handle free variables, constants, and applications. Attempted to convert: {e}" | _ => -- For other expression types (literals, lambdas, etc.), generate a placeholder name throwError m!"exprToConstructorExpr can only handle free variables, constants, and applications. Attempted to convert: {e}" /-- Classifies the arguments of an application `e`, one per argument position. An argument in an implicit or instance-implicit position that fails to classify is replaced by a `.Hole` instead of raising — such arguments are always re-inferable from the explicit ones. Arguments in explicit positions are classified normally and a failure there is a genuine error (propagated). -/ partial def classifyAppArgs (e : Expr) : MetaM (List ConstructorExpr) := do let fn := e.getAppFn let args := e.getAppArgs let fnType ← Meta.inferType fn Meta.forallBoundedTelescope fnType args.size fun bvars _ => do let mut result := #[] for h : i in [:args.size] do let isExplicit ← if h' : i < bvars.size then pure (← bvars[i].fvarId!.getDecl).binderInfo.isExplicit else pure true -- over-application: extra args are explicit if isExplicit then result := result.push (← exprToConstructorExpr args[i]!) else -- Implicit/instance position: try to classify, fall back to a hole. let ce ← try exprToConstructorExpr args[i]! catch _ => pure .Hole result := result.push ce return result.toList end /-- Converts an `Expr` to a `HypothesisExpr` if it is a variable or application of type constructor or function to constructor expressions else throws. We also pass in the current constructor name, as a poor person's location info. -/ def exprToHypothesisExpr (ctor : Name) (e : Expr) : MetaM HypothesisExpr := do trace[plausible.deriving.arbitrary] m!"Converting {e} to Hypexpr" let e ← Lean.Meta.withTransparency .reducible <| Meta.whnf e trace[plausible.deriving.arbitrary] m!"Converting {e} to Hypexpr after whnf" if e.isApp || e.isConst then let (ctorName, args) := e.getAppFnArgs let env ← getEnv if env.isConstructor ctorName then throwError m!"exprToHypothesisExpr: In constructor {ctor}\nExpr {e} cannot have head term {ctorName} which is a constructor. Must be a function or inductive" let constructorArgs ← args.mapM exprToConstructorExpr return (ctorName, constructorArgs.toList) else if e.isFVar then let name ← e.fvarId!.getUserName return (name, []) else match e with | .sort _lvl => return (`sort, []) | .proj structName idx expr => -- Structure field projections (e.g., T.Meta) — convert to application form let env ← getEnv let some info := getStructureInfo? env structName | throwError m!"exprToHypothesisExpr: {structName} is not a structure" let fieldName := structName ++ info.fieldNames[idx]! let constructorArg ← exprToConstructorExpr expr return (fieldName, [constructorArg]) | _ => throwError m!"exprToHypothesisExpr: In constructor {ctor} Expression\n{e}\nmust be of the form C a1 a2 ... an when used as hypothesis" /-- Helper function called by `updateSource`, which updates variables in a hypothesis `hyp` with the result of unification (provided via the `UnifyM` monad) -/ def updateNonRecSource (k : UnknownMap) (hyp : HypothesisExpr) : UnifyM Source := do let (ctorName, args) := hyp let updatedName ← UnifyM.findCanonicalUnknown k ctorName let updatedArgs ← List.mapM (UnifyM.updateConstructorArg k) args return .NonRec (updatedName, updatedArgs) /-- Updates a `Source` with the result of unification as contained in the `UnknownMap` -/ def updateSource (k : UnknownMap) (src : Source) : UnifyM Source := do match src with | .NonRec hyp => do updateNonRecSource k hyp | .Rec r tys => do let updatedTys ← List.mapM (UnifyM.updateConstructorArg k) tys return .Rec r updatedTys | .MutRec r tys => do let updatedTys ← List.mapM (UnifyM.updateConstructorArg k) tys return .MutRec r updatedTys /-- Updates a list of `ScheduleSteps` with the result of unification -/ def updateScheduleSteps (scheduleSteps : List ScheduleStep) : UnifyM (List ScheduleStep) := do UnifyM.withConstraints $ fun k => scheduleSteps.mapM (fun step => match step with | .Match e u p => do let updatedScrutinee ← UnifyM.findCanonicalUnknown k u let updatedPattern ← UnifyM.updatePattern k p return .Match e updatedScrutinee updatedPattern | .Unconstrained u src producerSort => do let updatedUnknown ← UnifyM.findCanonicalUnknown k u let updatedSrc ← updateSource k src return .Unconstrained updatedUnknown updatedSrc producerSort | .SuchThat unknownsAndTypes src dst => do let updatedUnknownsAndTypes ← unknownsAndTypes.mapM (fun (u, ty) => do let u' ← UnifyM.findCanonicalUnknown k u let ty' ← ty.mapM <| UnifyM.updateConstructorArg k return (u', ty')) let updatedSource ← updateSource k src return .SuchThat updatedUnknownsAndTypes updatedSource dst | .Check src polarity => do let updatedSrc ← updateSource k src return .Check updatedSrc polarity) /-- Takes the `patterns` and `equalities` fields from `UnifyState` (which are created after the conclusion of a constructor has been unified with the top-level arguments to the inductive relation), converts them to the appropriate `ScheduleStep`s, and prepends them to the `currentSchedule`. - The intuition for prepending these newly-created steps to the existing schedule is that we want to make sure all the equalities & pattern-matches needed for the conclusion hold before processing the rest of the schedule. -/ def addConclusionPatternsAndEqualitiesToSchedule (patterns : List (Unknown × Pattern)) (equalities : Std.HashSet (Unknown × Unknown)) (currentSchedule : Schedule) : Schedule := let (existingScheduleSteps, scheduleSort) := currentSchedule let matchSteps := (Function.uncurry (ScheduleStep.Match .allowImplicit)) <$> patterns -- We should never have an equality here. Assertion after unification should handle that though. let equalityCheckSteps := (fun (u1, u2) => ScheduleStep.Check (Source.NonRec (``Eq, [.Unknown u1, .Unknown u2])) true) <$> equalities.toList (matchSteps ++ equalityCheckSteps ++ existingScheduleSteps, scheduleSort) /-- Extracts all variable names from a ConstructorExpr (looks inside constructors) -/ partial def varsInConstructorExpr : ConstructorExpr → List Name | .Unknown u => [u] | .Ctor _ args | .FuncApp _ args | .TyCtor _ args => args.flatMap varsInConstructorExpr | .Lit _ | .CSort _ | .Hole => [] /-- Computes output indices: position i is an output if any output variable appears in it -/ def computeOutputIndicesForRewrite (hypArgs : List ConstructorExpr) (outputVarNames : List Name) : List Nat := filterMapWithIndex (fun i arg => let vars := varsInConstructorExpr arg if vars.any (· ∈ outputVarNames) then some i else none) hypArgs /-- Rewrites `Source.NonRec` calls in schedule steps to `Source.MutRec` when the hypothesis matches a sibling spec exactly (same inductive, same output positions, same derive sort). `siblings` is `(inductiveName, outputIndices, auxFnName, siblingDeriveSort)`. -/ def rewriteMutualCalls (steps : List ScheduleStep) (siblings : List (Name × List Nat × Name × DeriveSort)) : List ScheduleStep := let matchesSibling (hyp : HypothesisExpr) (outNames : List Name) (stepDeriveSort : DeriveSort) : Option (Name × List Nat) := let (hypName, hypArgs) := hyp let stepOutputIdxs := computeOutputIndicesForRewrite hypArgs outNames siblings.findSome? fun (indName, sibOutputIdxs, auxName, sibSort) => if hypName == indName && stepOutputIdxs == sibOutputIdxs && stepDeriveSort == sibSort then some (auxName, sibOutputIdxs) else none steps.map fun step => match step with | .SuchThat vs (.NonRec hyp) ps => let ds := match ps with | .Generator => DeriveSort.Generator | .Enumerator => .Enumerator match matchesSibling hyp (vs.map Prod.fst) ds with | some (auxName, outputIdxs) => let (_, hypArgs) := hyp let inputArgs := filterWithIndex (fun i _ => i ∉ outputIdxs) hypArgs .SuchThat vs (.MutRec auxName inputArgs) ps | none => step | .Unconstrained v (.NonRec hyp) ps => let ds := match ps with | .Generator => DeriveSort.Generator | .Enumerator => .Enumerator match matchesSibling hyp [v] ds with | some (auxName, outputIdxs) => let (_, hypArgs) := hyp let inputArgs := filterWithIndex (fun i _ => i ∉ outputIdxs) hypArgs .Unconstrained v (.MutRec auxName inputArgs) ps | none => step | .Check (.NonRec hyp) pol => match matchesSibling hyp [] .Checker with | some (auxName, _) => let (_, hypArgs) := hyp .Check (.MutRec auxName hypArgs) pol | none => step | other => other /-- Identifies a unique derivation spec (inductive + output mode + sort). -/ structure SpecKey where inductiveName : Name outputIndices : List Nat deriveSort : DeriveSort deriving Repr, BEq, Hashable, Inhabited /-- Pretty-prints a SpecKey as a spec form like "fun τ => ∃ Γ e, typing Γ e τ". Requires knowing the inductive's arg types (number of args). -/ def SpecKey.prettyPrint (key : SpecKey) (numArgs : Nat) : String := let varPool := #["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"] let (inputs, outputs, appArgs) := Id.run do let mut inputs : List String := [] let mut outputs : List String := [] let mut appArgs : List String := [] for i in List.range numArgs do let name := varPool.getD i s!"v{i}" if i ∈ key.outputIndices then outputs := outputs ++ [name] else inputs := inputs ++ [name] appArgs := appArgs ++ [name] (inputs, outputs, appArgs) let inputsStr := if inputs.isEmpty then "" else s!"fun {String.intercalate " " inputs} => " let outputsStr := if outputs.isEmpty then "" else s!"∃ {String.intercalate " " outputs}, " let sortStr := match key.deriveSort with | .Generator => "[generator] " | .Enumerator => "[enumerator] " | .Checker => "[checker] " | .Theorem => "[theorem] " s!"{sortStr}{inputsStr}{outputsStr}{key.inductiveName} {String.intercalate " " appArgs}" /-- Score for a derived schedule (used for selecting best schedule). -/ structure SpecScore where checks : Nat := 0 unconstrained : Nat := 0 backtracking : Nat := 0 deriving Repr, BEq, Ord instance : Inhabited SpecScore := ⟨{}⟩ /-- A complete derivation plan for one inductive at one output mode. Analogous to QuickChick's `inductive_schedule`. -/ structure InductiveSchedule where /-- The spec this schedule is for -/ key : SpecKey /-- Freshened argument names (used in the schedule steps) -/ argNames : List Name /-- The recursive function name used in Source.Rec calls -/ recFnName : Name /-- Per-constructor schedules for non-recursive (base) constructors -/ baseSchedules : List (Name × Schedule) /-- Per-constructor schedules for recursive constructors -/ recSchedules : List (Name × Schedule) /-- Quality score for this derivation (type-erased, from active scoring bundle) -/ score : Score /-- True if this spec already has an instance in the environment (no need to compile) -/ alreadyExists : Bool := false /-- Time taken to derive the full spec (in microseconds, includes dep derivation) -/ derivationTimeUs : Nat := 0 /-- Per-constructor stats: (name, time in μs, schedules considered, score) -/ ctorStats : List (Name × Nat × Nat × Score) := [] deriving Repr /-- Result of deriving a schedule, stored in the dependency memo. -/ inductive MemoEntry | inProgress -- derivation started, cycle detection | done (indSched : InductiveSchedule) | failed (msg : String) deriving Repr /-- Whether a dependency is for a base type (needs Arbitrary/Enum) or a relation (needs ArbitrarySizedSuchThat/EnumSizedSuchThat/DecOpt). -/ inductive DepKind /-- Unconstrained generation of a base type (needs `Arbitrary` or `Enum`) -/ | baseType /-- Constrained generation from a relation (needs `ArbitrarySizedSuchThat` or `EnumSizedSuchThat`) -/ | relation /-- Checking a relation (needs `DecOpt`) -/ | checker deriving Repr, BEq /-- A dependency extracted from a schedule: what instance is needed. - `kind`: whether this is a base type, constrained relation, or checker - `inductiveName`: the type/relation being referenced - `hypothesis`: the full hypothesis expression (indName + args) - `outputVarNames`: the variables being produced (determines output positions) - `outputIndices`: positions in the hypothesis args that are outputs - `deriveSort`: the producer sort context (Generator or Enumerator) -/ structure ScheduleDep where kind : DepKind inductiveName : Name hypothesis : HypothesisExpr outputVarNames : List Name outputIndices : List Nat deriveSort : DeriveSort deriving Repr, BEq /-- Computes output indices: which positions in hypArgs are output variables -/ def computeOutputIndices (hypArgs : List ConstructorExpr) (outputVarNames : List Name) : List Nat := filterMapWithIndex (fun i arg => match arg with | .Unknown name => if name ∈ outputVarNames then some i else none | _ => none) hypArgs /-- Extracts all `Source.NonRec` dependencies from schedule steps. Each dependency tells us what instance is needed: which inductive, which argument positions are outputs, and what derive sort. -/ def collectNonRecDeps (steps : List ScheduleStep) : List ScheduleDep := steps.filterMap fun step => match step with | .SuchThat vs (.NonRec hyp) ps => let ds := match ps with | .Generator => DeriveSort.Generator | .Enumerator => .Enumerator let outNames := vs.map Prod.fst some { kind := .relation inductiveName := hyp.fst hypothesis := hyp outputVarNames := outNames outputIndices := computeOutputIndices hyp.snd outNames deriveSort := ds } | .Unconstrained v (.NonRec hyp) ps => let ds := match ps with | .Generator => DeriveSort.Generator | .Enumerator => .Enumerator some { kind := .baseType inductiveName := hyp.fst hypothesis := hyp outputVarNames := [v] outputIndices := computeOutputIndices hyp.snd [v] deriveSort := ds } | .Check (.NonRec hyp) _ => some { kind := .checker inductiveName := hyp.fst hypothesis := hyp outputVarNames := [] outputIndices := [] deriveSort := .Checker } | _ => none /-- DFS from a root SpecKey through chosen schedules to find all actually-used dependencies. -/ partial def collectUsedDeps (root : SpecKey) (memo : Std.HashMap SpecKey MemoEntry) (visited : Std.HashSet SpecKey := {}) : Std.HashSet SpecKey := if visited.contains root then visited else let visited := visited.insert root match memo[root]? with | some (.done indSched) => let allSchedules := indSched.baseSchedules ++ indSched.recSchedules let deps := allSchedules.flatMap (fun (_, (steps, _)) => collectNonRecDeps steps) let relDeps := deps.filter (fun d => d.kind == .relation || d.kind == .checker) relDeps.foldl (fun acc dep => let depKey : SpecKey := { inductiveName := dep.inductiveName, outputIndices := dep.outputIndices, deriveSort := dep.deriveSort } collectUsedDeps depKey memo acc) visited | _ => visited /-- Given a set of used SpecKeys and the memo, compute SCCs (mutual groups). Returns components in topological order (dependencies before dependants). -/ def computeSpecSCC (usedKeys : List SpecKey) (memo : Std.HashMap SpecKey MemoEntry) : List (List SpecKey) := let successors (key : SpecKey) : List SpecKey := match memo[key]? with | some (.done indSched) => let allScheds := indSched.baseSchedules ++ indSched.recSchedules let deps := allScheds.flatMap (fun (_, (steps, _)) => collectNonRecDeps steps) let relDeps := deps.filter (fun d => d.kind == .relation || d.kind == .checker) let depKeys := relDeps.map (fun d => SpecKey.mk d.inductiveName d.outputIndices d.deriveSort) depKeys.filter (usedKeys.contains ·) | _ => [] Lean.SCC.scc usedKeys successors /-- Checks if any step in a schedule uses `Source.MutRec`. -/ def scheduleUsesMutualCall (steps : List ScheduleStep) : Bool := steps.any fun step => match step with | .Unconstrained _ (.MutRec ..) _ => true | .SuchThat _ (.MutRec ..) _ => true | .Check (.MutRec ..) _ => true | _ => false end Schedules