refactor: generalize buildRBTop

Lake/BuildModule.lean

@@ -155,13 +155,13 @@ def  Package.recFetchModuleOleanTargetWithLocalImports [Monad m] [MonadLiftT Bui
 
 -- ## Definitions
 
-abbrev ModuleFetchM (α) :=
+abbrev ModuleBuildM (α) :=
   -- equivalent to `RBTopT (cmp := Name.quickCmp) Name α BuildM`.
   -- phrased this way to use `NameMap`
   EStateT (List Name) (NameMap α) BuildM
 
 abbrev ModuleFetch (α) :=
-  RecFetch Name α (ModuleFetchM α)
+  RecFetch Name α (ModuleBuildM α)
 
 abbrev OleanTargetFetch := ModuleFetch ActiveFileTarget
 abbrev OleanAndCTargetFetch := ModuleFetch ActiveOleanAndCTarget

Lake/BuildPackage.lean

@@ -39,14 +39,16 @@ def Package.buildModuleOleanAndCTargetDAG
 (self : Package) : BuildM (Array ActiveOleanAndCTarget × OleanAndCTargetMap) := do
   let fetch : OleanAndCTargetFetch :=
     self.recFetchModuleOleanAndCTargetWithLocalImports moreOleanDirs depTarget
-  failOnImportCycle <| ← mods.mapM (buildRBTop fetch) |>.run {}
+  let (resE, map) ← mods.mapM (buildRBTop fetch) |>.run {}
+  (← failOnImportCycle resE, map)
 
 def Package.buildModuleOleanTargetDAG
 (mods : Array Name) (moreOleanDirs : List FilePath) (depTarget : ActiveBuildTarget x)
 (self : Package) : BuildM (Array ActiveFileTarget × OleanTargetMap) := do
   let fetch : OleanTargetFetch :=
     self.recFetchModuleOleanTargetWithLocalImports moreOleanDirs depTarget
-  failOnImportCycle <| ← mods.mapM (buildRBTop fetch) |>.run {}
+  let (resE, map) ← mods.mapM (buildRBTop fetch) |>.run {}
+  (← failOnImportCycle resE, map)
 
 def Package.buildOleanAndCTargetDAG
 (moreOleanDirs : List FilePath) (depTarget : ActiveBuildTarget x)

Lake/BuildTop.lean

@@ -1,26 +1,83 @@
 /-
-Copyright (c) 2017 Microsoft Corporation. All rights reserved.
+Copyright (c) 2021 Microsoft Corporation. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Gabriel Ebner, Sebastian Ullrich, Mac Malone
+Authors: Sebastian Ullrich, Mac Malone
 -/
 import Std.Data.RBMap
 
 open Std
 namespace Lake
 
+--------------------------------------------------------------------------------
+-- # Topological / Suspending Builder
+--------------------------------------------------------------------------------
+
+-- ## Abstractions
+
 /-- A recursive object fetcher. -/
 def RecFetch.{u,v,w} (k : Type u) (o : Type v) (m : Type v → Type w) :=
   k → (k → m o) → m o
 
-/-- A exception plus state monad transformer (i.e., `StateT` + `ExceptT`). -/
-abbrev EStateT (ε σ m) :=
-  StateT σ <| ExceptT ε m
+/-- A monad equipped with a key-object store. -/
+class MonadStore (k : Type u) (o : Type v) (m : Type v → Type w) where
+  fetch? : k → m (Option o)
+  store : k → o → m PUnit
 
-def EStateT.run (state : σ) (self : EStateT ε σ m α) : m (Except ε (α × σ)) :=
-  StateT.run self state |>.run
+export MonadStore (fetch? store)
+
+instance [Monad m] [MonadStore k o m] : MonadStore k o (ExceptT ε m) where
+  fetch? key := liftM (m := m) <| fetch? key
+  store key obj := liftM (m := m) <| store key obj
+
+-- ## Algorithm
+
+/-- Auxiliary function for `buildTop`. -/
+partial def buildTopCore
+{k o} {m} [BEq k] [Inhabited o] [Monad m] [MonadStore k o m]
+(parents : List k) (build : RecFetch k o (ExceptT (List k) m)) (key : k) : ExceptT (List k) m o := do
+  -- detect cyclic builds
+  if parents.contains key then
+    throw <| key :: (parents.partition (· != key)).1 ++ [key]
+  -- return previous output if already built
+  if let some output ← fetch? key then
+    return output
+  -- build the key recursively
+  let output ← build key (buildTopCore (key :: parents) build)
+  -- save the output (to prevent repeated builds of the same key)
+  store key output
+  return output
+
+/--
+  Recursively fills a `MonadStore` of key-object pairs by
+  fetching objects topologically (i.e., via a depth-first search with memoization).
+  Called a suspending scheduler in "Build systems à la carte".
+-/
+def buildTop
+{k o} {m} [BEq k] [Inhabited o] [Monad m] [MonadStore k o m]
+(build : RecFetch k o (ExceptT (List k) m)) (key : k) : ExceptT (List k) m o :=
+  buildTopCore [] build key
+
+--------------------------------------------------------------------------------
+-- # RBMap Version
+--------------------------------------------------------------------------------
+
+/-- A exception plus state monad transformer (i.e., `ExceptT` + `StateT`). -/
+abbrev EStateT (ε σ m) :=
+  ExceptT ε <| StateT σ m
+
+def EStateT.run (state : σ) (self : EStateT ε σ m α) : m (Except ε α × σ) :=
+  ExceptT.run self |>.run state
 
 def EStateT.run' [Monad m] (state : σ) (self : EStateT ε σ m α) : m (Except ε α) :=
-  StateT.run' self state |>.run
+  ExceptT.run self |>.run' state
+
+/-- A transformer that adds an `RBMap` store to a monad. -/
+abbrev RBMapT.{u,v} (k : Type u) (o : Type u) (cmp) (m : Type u → Type v) :=
+  StateT (RBMap k o cmp) m
+
+instance (cmp) [Monad m] : MonadStore k o (RBMapT k o cmp m) where
+  fetch? key := do (← get).find? key
+  store key obj := modify (·.insert key obj)
 
 /--
   Monad transformer for an RBMap-based topological walk.
@@ -29,28 +86,7 @@ def EStateT.run' [Monad m] (state : σ) (self : EStateT ε σ m α) : m (Except
 abbrev RBTopT.{u,v} (k : Type u) (o : Type u) (cmp) (m : Type u → Type v) :=
   EStateT (List k) (RBMap k o cmp) m
 
-/-- Auxiliary function for `buildRBTop`. -/
-partial def buildRBTopCore
-{k o} {cmp} {m : Type u → Type u} [BEq k] [Inhabited o] [Monad m]
-(parents : List k) (fetch : RecFetch k o (RBTopT k o cmp m)) (key : k)
-: RBTopT k o cmp m o := do
-  -- detect cyclic builds
-  if parents.contains key then
-    throw <| key :: (parents.partition (· != key)).1 ++ [key]
-  -- return previous output if already built
-  if let some output := (← get).find? key then
-    return output
-  -- build the key recursively
-  let output ← fetch key (buildRBTopCore (key :: parents) fetch)
-  -- save the output (to prevent repeated builds of the same key)
-  modify (·.insert key output)
-  return output
-
-/--
-  Recursively constructs an `RBMap` of key-object pairs by
-  fetching objects topologically (i.e., via a depth-first search with memoization).
-  Called a suspending scheduler in "Build systems à la carte".
--/
+/-- The `RBMap` version of `buildTop`. -/
 def buildRBTop {k o} {cmp} {m} [BEq k] [Inhabited o] [Monad m]
-(fetch : RecFetch k o (RBTopT k o cmp m)) (key : k) : RBTopT k o cmp m o :=
-  buildRBTopCore [] fetch key
+(build : RecFetch k o (RBTopT k o cmp m)) (key : k) : RBTopT k o cmp m o :=
+  buildTop build key