refactor: lake: fix setDepPkgs & further verify resolution (#13487)

This PR fixes an issue with `Workspace.setDepPkgs` (introduced in #13445) where it did not properly update the workspace `packageMap`. In addition, this PR further verifies aspects of the workspace construction during dependency resolution.
2026-04-20 23:47:27 -04:00 · 2026-04-20 23:47:27 -04:00 · fcaebdad22
commit fcaebdad22
parent e2f9df6578
2 changed files with 130 additions and 90 deletions
--- a/src/lake/Lake/Config/Workspace.lean
+++ b/src/lake/Lake/Config/Workspace.lean
@ -14,6 +14,8 @@ public import Lake.Config.TargetConfig
 public import Lake.Config.LakeConfig
 meta import Lake.Util.OpaqueType
 import Lean.DocString.Syntax
+import Init.Data.Range.Polymorphic.Iterators
+import Init.Data.Range.Polymorphic.Lemmas

 set_option doc.verso true

@ -63,18 +65,25 @@ public structure Workspace.Raw.WF (ws : Workspace.Raw) : Prop where
 /-- A Lake workspace -- the top-level package directory. -/
 public structure Workspace extends raw : Workspace.Raw, wf : raw.WF

-public instance : Nonempty Workspace := .intro {
+/-- Constructs an arbitrary well-formed workspace with {lean}`n` packages. -/
+noncomputable def Workspace.ofSize (n : Nat) (h : 0 < n) : Workspace := {
  lakeEnv := default
  lakeConfig := Classical.ofNonempty
  lakeCache := Classical.ofNonempty
-  packages := #[{(Classical.ofNonempty : Package) with wsIdx := 0}]
-  size_packages_pos := by simp
+  packages := (0...<n).toArray.map fun i =>
+    {(Classical.ofNonempty : Package) with wsIdx := i}
+  size_packages_pos := by
+    simp [Std.Rco.size, Std.Rxo.HasSize.size, Std.Rxc.HasSize.size, h]
  packages_wsIdx {i} h := by
-    cases i with
-    | zero => simp
-    | succ => simp at h
+    simp [Std.Rco.getElem_toArray_eq, Std.PRange.succMany?]
 }

+theorem Workspace.size_packages_ofSize :
+  (ofSize n h).packages.size = n
+:= by simp [ofSize, Std.Rco.size, Std.Rxo.HasSize.size, Std.Rxc.HasSize.size]
+
+public instance : Nonempty Workspace := ⟨.ofSize 1 Nat.zero_lt_one⟩
+
 public hydrate_opaque_type OpaqueWorkspace Workspace

 /-- Returns the names of the root modules of the package's default targets. -/
@ -93,6 +102,11 @@ namespace Workspace
@[inline] public def root (self : Workspace) : Package :=
  self.packages[0]'self.size_packages_pos

+/-- **For internal use only.** -/
+public theorem wsIdx_root_lt {ws : Workspace} :
+  ws.root.wsIdx < ws.packages.size
+:= ws.packages_wsIdx _ ▸ ws.size_packages_pos
+
 /-- **For internal use.** Whether this workspace is Lean itself.  -/
@[inline] def bootstrap (self : Workspace) : Bool :=
  self.root.bootstrap
--- a/src/lake/Lake/Load/Resolve.lean
+++ b/src/lake/Lake/Load/Resolve.lean
@ -11,12 +11,12 @@ public import Lake.Load.Manifest
 import Lake.Util.IO
 import Lake.Util.StoreInsts
 import Lake.Config.Monad
-import Lake.Build.Topological
 import Lake.Load.Materialize
 import Lake.Load.Lean.Eval
 import Lake.Load.Package
 import Init.Data.Range.Polymorphic.Iterators
 import Init.TacticsExtra
+import Lean.Runtime

 open System Lean

@ -71,69 +71,41 @@ def Workspace.addDepPackage'
  let ws := ws.addPackage' pkg wsIdx_mkPackage |>.addFacetDecls fileCfg.facetDecls
  return ⟨ws, by simp [ws, packages_addFacetDecls, packages_addPackage']⟩

-/--
-The resolver's call stack of dependencies.
-That is, the dependency currently being resolved plus its parents.
-/
-abbrev DepStack := CallStack Name
-
-/--
-A monad transformer for recursive dependency resolution.
-It equips the monad with the stack of dependencies currently being resolved.
-/
-abbrev DepStackT m := CallStackT Name m
-
-@[inline] nonrec def DepStackT.run
-  (x : DepStackT m α) (stack : DepStack := {})
-: m α := x.run stack
-
-/-- Log dependency cycle and error. -/
-@[specialize] def depCycleError [MonadError m] (cycle : Cycle Name) : m α :=
-  error s!"dependency cycle detected:\n{formatCycle cycle}"
-
-instance [Monad m] [MonadError m] : MonadCycleOf Name (DepStackT m) where
-  throwCycle := depCycleError
-
-/-- The monad of the dependency resolver. -/
-abbrev ResolveT m := DepStackT <| StateT Workspace m
-
-@[inline] nonrec def ResolveT.run
-  (ws : Workspace) (x : ResolveT m α) (stack : DepStack := {})
-: m (α × Workspace) := x.run stack |>.run ws
-
-/-- Recursively run a `ResolveT` monad starting from the workspace's root. -/
-@[specialize] def Workspace.runResolveT
-  [Monad m] [MonadError m] (ws : Workspace)
-  (go : RecFetchFn Package PUnit (ResolveT m))
-  (root := ws.root) (stack : DepStack  := {})
-: m Workspace := do
-  let (_, ws) ← ResolveT.run ws (stack := stack) do
-    recFetchAcyclic (·.baseName) go root
-  return ws

 def Workspace.setDepPkgs
-  (self : Workspace) (wsIdx : Nat) (depPkgs : Array Package)
-: Workspace := {self with
-  packages := self.packages.modify wsIdx ({· with depPkgs})
-  size_packages_pos := by simp [self.size_packages_pos]
-  packages_wsIdx {i} := by
-    if h : wsIdx = i then
-      simp [h, Array.getElem_modify_self, self.packages_wsIdx]
-    else
-      simp [Array.getElem_modify_of_ne h, self.packages_wsIdx]
-}
+  (self : Workspace) (pkg : Package) (depPkgs : Array Package)
+  (h : pkg.wsIdx < self.packages.size)
+: Workspace :=
+  let pkg := {pkg with depPkgs}
+  {self with
+    packages := self.packages.set pkg.wsIdx pkg h
+    packageMap := self.packageMap.insert pkg.keyName pkg
+    size_packages_pos := by simp [self.size_packages_pos]
+    packages_wsIdx {i} := by
+      intro hi
+      rw [Array.size_set] at hi
+      rw [self.packages.getElem_set]
+      split
+      · assumption
+      · rw [self.packages_wsIdx]
+  }

-structure ResolveState where
+theorem Workspace.size_packages_setDepPkgs :
+  (setDepPkgs ws pkg depPkgs h).packages.size = ws.packages.size
+:= by simp [setDepPkgs]
+
+structure ResolveState (start : Nat) where
  ws : Workspace
  depIdxs : Array Nat
  lt_of_mem : ∀ i ∈ depIdxs, i < ws.packages.size
+  start_le : start ≤ ws.packages.size

 namespace ResolveState

-@[inline] def init (ws : Workspace) (size : Nat) : ResolveState :=
-  {ws, depIdxs := Array.mkEmpty size, lt_of_mem := by simp}
+@[inline] def init (ws : Workspace) (size : Nat) : ResolveState ws.packages.size :=
+  {ws, depIdxs := Array.mkEmpty size, lt_of_mem := by simp, start_le := Nat.le_refl ..}

-@[inline] def reuseDep (s : ResolveState) (wsIdx : Fin s.ws.packages.size) : ResolveState :=
+@[inline] def reuseDep (s : ResolveState n) (wsIdx : Fin s.ws.packages.size) : ResolveState n :=
  have lt_of_mem := by
    intro i i_mem
    cases Array.mem_push.mp i_mem with
@ -142,10 +114,10 @@ namespace ResolveState
  {s with depIdxs := s.depIdxs.push wsIdx, lt_of_mem}

@[inline] def newDep
-  (s : ResolveState) (dep : MaterializedDep)
+  (s : ResolveState n) (dep : MaterializedDep)
  (lakeOpts : NameMap String) (leanOpts : Options) (reconfigure : Bool)
-: LogIO ResolveState := do
-  let ⟨ws, depIdxs, lt_of_mem⟩ := s
+: LogIO (ResolveState n) := do
+  let {ws, depIdxs, lt_of_mem, start_le} := s
  let wsIdx := ws.packages.size
  let ⟨ws', h⟩ ← ws.addDepPackage' dep lakeOpts leanOpts reconfigure
  have lt_of_mem := by
@ -153,10 +125,31 @@ namespace ResolveState
    cases Array.mem_push.mp i_mem with
    | inl i_mem => exact h ▸ Nat.lt_add_one_of_lt (lt_of_mem i i_mem)
    | inr i_eq => simp only [wsIdx, i_eq, h, Nat.lt_add_one]
-  return ⟨ws', depIdxs.push wsIdx, lt_of_mem⟩
+  have start_le := Nat.le_trans start_le <| h ▸ Nat.le_add_right ..
+  return ⟨ws', depIdxs.push wsIdx, lt_of_mem, start_le⟩

 end ResolveState

+abbrev MinWorkspace (ws : Workspace) :=
+  {ws' : Workspace // ws.packages.size ≤ ws'.packages.size}
+
+instance : Nonempty (MinWorkspace ws) := ⟨⟨ws, Nat.le_refl ..⟩⟩
+
+@[inline] unsafe def guardBySizeImpl [Pure m] [MonadError m] (as : Array α) : m (PLift (as.size ≤ Lean.maxSmallNat)) :=
+  pure ⟨lcProof⟩
+
+/--
+Returns a proof that the size of an `Array` is at most `Lean.maxSmallNat`.
+
+This is modelled to fail via `MonadError` if this property does not hold. However, when compiled,
+this is implemented by a no-op, because this is a fixed property of the Lean runtime.
+
+This function can be used to prove that Array-bounded recursion terminates.
+-/
+@[implemented_by guardBySizeImpl]
+def guardBySize! [Pure m] [MonadError m] (as : Array α) : m (PLift (as.size ≤ Lean.maxSmallNat)) :=
+  if h : as.size ≤ Lean.maxSmallNat then pure ⟨h⟩ else error "Array-bounded termination"
+
 /-
 Recursively visits each node in a package's dependency graph, starting from
 the workspace package `root`. Each dependency missing from the workspace is
@ -170,36 +163,51 @@ See `Workspace.updateAndMaterializeCore` for more details.
@[inline] def Workspace.resolveDepsCore
  [Monad m] [MonadError m] [MonadLiftT LogIO m] (ws : Workspace)
  (resolve : Package → Dependency → Workspace → m MaterializedDep)
-  (root : Package := ws.root) (stack : DepStack := {})
+  (root : Nat := 0) (h : root < ws.packages.size := by exact ws.size_packages_pos)
  (leanOpts : Options := {}) (reconfigure := true)
-: m Workspace := do
-  ws.runResolveT go root stack
+: m (MinWorkspace ws) := do
+  go ws root h
 where
-  @[specialize] go pkg recurse : ResolveT m Unit := fun depStack ws => do
+  @[specialize] go
+    (ws : Workspace) (wsIdx : Nat) (h : wsIdx < ws.packages.size)
+  : m (MinWorkspace ws) := do
    let start := ws.packages.size
+    let pkg : Package := ws.packages[wsIdx]
+    have lt_start : pkg.wsIdx < start := ws.packages_wsIdx _ ▸ h
    -- Materialize and load the missing direct dependencies of `pkg`
    let s := ResolveState.init ws pkg.depConfigs.size
-    let ⟨ws, depIdxs, lt_of_mem⟩ ← pkg.depConfigs.foldrM (m := m) (init := s) fun dep s => do
+    let ⟨ws, depIdxs, lt_of_mem, start_le⟩ ← pkg.depConfigs.foldrM (m := m) (init := s) fun dep s => do
      if let some wsIdx := s.ws.packages.findFinIdx? (·.baseName == dep.name) then
        return s.reuseDep wsIdx -- already handled in another branch
      if pkg.baseName = dep.name then
        error s!"{pkg.prettyName}: package requires itself (or a package with the same name)"
      let matDep ← resolve pkg dep s.ws
      s.newDep matDep dep.opts leanOpts reconfigure
-    let depsEnd := ws.packages.size
    -- Recursively load the dependencies' dependencies
-    let ws ← ws.packages.foldlM (start := start) (init := ws) fun ws pkg =>
-      (·.2) <$> recurse pkg depStack ws
+    let stop := ws.packages.size
+    let ⟨le_maxSmallNat⟩ ← guardBySize! ws.packages
+    let ⟨ws, stop_le⟩ ← id do
+      let mut ws' : {ws' : Workspace // stop ≤ ws'.packages.size} := ⟨ws, Nat.le_refl _⟩
+      for h_rco : wsIdx in (start...<stop) do
+        let ⟨ws, stop_le⟩ := ws'
+        have lt_stop := Std.Rco.lt_upper_of_mem h_rco
+        let ⟨ws, h⟩ ← go ws wsIdx <| Nat.lt_of_lt_of_le lt_stop stop_le
+        ws' := ⟨ws, Nat.le_trans stop_le h⟩
+      return ws'
+    have start_le := Nat.le_trans start_le stop_le
    -- Add the package's dependencies to the package
-    let ws :=
-      if h_le : depsEnd ≤ ws.packages.size then
-        ws.setDepPkgs pkg.wsIdx <| depIdxs.attach.map fun ⟨wsIdx, h_mem⟩ =>
-          ws.packages[wsIdx]'(Nat.lt_of_lt_of_le (lt_of_mem wsIdx h_mem) h_le)
-      else
-        have : Inhabited Workspace := ⟨ws⟩
-        panic! "workspace shrunk" -- should be unreachable
-    return ((), ws)
-
+    let depPkgs := depIdxs.attach.map fun ⟨wsIdx, h_mem⟩ =>
+      ws.packages[wsIdx]'(Nat.lt_of_lt_of_le (lt_of_mem wsIdx h_mem) stop_le)
+    let ws := ws.setDepPkgs pkg depPkgs <| Nat.lt_of_lt_of_le lt_start start_le
+    have start_le := Nat.le_trans start_le <| by
+      simp [ws, Workspace.size_packages_setDepPkgs]
+    return ⟨ws, start_le⟩
+  termination_by Lean.maxSmallNat - wsIdx
+  decreasing_by
+    have start_le := Std.Rco.lower_le_of_mem h_rco
+    have lt_wsIdx := Nat.lt_of_lt_of_le h start_le
+    refine Nat.sub_lt_sub_left (Nat.lt_of_lt_of_le lt_wsIdx ?_) lt_wsIdx
+    exact Nat.le_trans (Nat.le_of_lt lt_stop) le_maxSmallNat

 /--
 Adds monad state used to update the manifest.
@ -430,22 +438,40 @@ def Workspace.updateAndMaterializeCore
 : LoggerIO (Workspace × NameMap PackageEntry) := UpdateT.run do
  reuseManifest ws toUpdate
  if updateToolchain then
-    let deps := ws.root.depConfigs.reverse
+    let numDeps := ws.root.depConfigs.size
+    -- Update and materialize the top-level dependenciess
+    let deps : Vector _ numDeps := Vector.mk ws.root.depConfigs.reverse (by simp [numDeps])
    let matDeps ← deps.mapM fun dep => do
      logVerbose s!"{ws.root.prettyName}: updating '{dep.name}' with {toJson dep.opts}"
      updateAndMaterializeDep ws ws.root dep
-    ws.updateToolchain matDeps
+    -- Update the toolchain based on the top-level dependenciess
+    ws.updateToolchain matDeps.toArray
+    -- Load the top-level dependenciess
    let start := ws.packages.size
-    let ws ← (deps.zip matDeps).foldlM (init := ws) fun ws (dep, matDep) => do
-      addDependencyEntries matDep
-      let ⟨ws, _⟩ ← ws.addDepPackage' matDep dep.opts leanOpts true
-      return ws
+    let ⟨ws, h⟩ ← id do
+      let mut ws' : {ws : Workspace // start ≤ ws.packages.size} := ⟨ws, Nat.le_refl _⟩
+      for h : i in 0...<numDeps do
+        let matDep := matDeps[i]
+        addDependencyEntries matDep
+        let lakeOpts := deps[i].opts
+        let ⟨ws, h⟩ ← ws'.val.addDepPackage' matDep lakeOpts leanOpts true
+        ws' := ⟨ws, Nat.le_trans ws'.property <| by simp [h]⟩
+      return ws'
    let stop := ws.packages.size
-    let ws ← ws.packages.foldlM (init := ws) (start := start) fun ws pkg =>
-      ws.resolveDepsCore updateAndAddDep pkg [ws.root.baseName] leanOpts true
+    have start_le_stop : start ≤ stop := h
+    -- Resolve the top-level dependencies' dependencies'
+    let ⟨ws, _⟩ ← id do
+      let mut ws' : {ws : Workspace // stop ≤ ws.packages.size} := ⟨ws, Nat.le_refl _⟩
+      for h : wsIdx in start...<stop do
+        let ⟨ws, stop_le⟩ := ws'
+        have h := Nat.lt_of_lt_of_le (Std.Rco.lt_upper_of_mem h) stop_le
+        let ⟨ws, h⟩ ← ws.resolveDepsCore updateAndAddDep wsIdx h
+          (leanOpts := leanOpts) (reconfigure := true)
+        ws' := ⟨ws, Nat.le_trans stop_le h⟩
+      return ws'
    -- Set dependency packages after `resolveDepsCore` so
    -- that the dependencies' dependencies are also properly set
-    return ws.setDepPkgs ws.root.wsIdx ws.packages[start...<stop]
+    return ws.setDepPkgs ws.root ws.packages[start...<stop] ws.wsIdx_root_lt
  else
    ws.resolveDepsCore updateAndAddDep (leanOpts := leanOpts) (reconfigure := true)
 where