chore: basic tests for DeclNeeds and ImportNeeds

ImportGraph/Shake/DeclNeeds.lean

@@ -337,6 +337,14 @@ inductive ComptimeDependency where
   stx (kind : SyntaxNodeKind) (pos? : Option Syntax.Range) -- TODO: just record syntax?
 deriving Inhabited, BEq, Repr, Hashable, Ord
 
+def ComptimeDependency.toMessageData : ComptimeDependency → MessageData
+  | .indirect { kind, declName } mods =>
+    m!"indirect({kind} for {.ofConstName declName} from {mods})"
+  | .stx kind pos? => m!"syntax({.ofConstName kind (fullNames := true)}\
+    {if let some pos := pos? then m!"@⟨{pos.start},{pos.stop}⟩" else ""})"
+
+instance : ToMessageData ComptimeDependency := ⟨(·.toMessageData)⟩
+
 /-- A location at which a `ConstantInfo` can reference another. -/
 inductive ConstLocation where | type | value
 deriving Inhabited, BEq, Repr, Hashable, Ord
@@ -392,15 +400,17 @@ deriving Inhabited, BEq, Repr, Hashable, Ord
 namespace DeclDeclNeedsKind
 
 /-- An English-language description of a `DeclDeclNeedsKind`. -/
-def pretty : DeclDeclNeedsKind → String
-  | .expr down reExported up => s!"uses {up} in {down}{
+def toMessageData : DeclDeclNeedsKind → MessageData
+  | .expr down reExported up => m!"{if up matches .type then "used" else s!"uses {up}"} in {down}{
     match reExported with
     | some true => " (public)"
     | some false => " (private)"
     | none => ""}"
   | .runtimeIR => "used in runtime IR"
   | .metaIR _ _ => "used in meta IR" -- TODO
-  | .comptime _ => "comptime dependency" -- TODO
+  | .comptime source => m!"{source}"
+
+instance : ToMessageData DeclDeclNeedsKind := ⟨(·.toMessageData)⟩
 
 /-- `some .comptime` if this demands the source declaration's `IRPhases` includes `.comptime`;
 `some .runtime` if this demands the source declaration's `IRPhases` includes `.runtime`. (`.all`
@@ -440,6 +450,19 @@ structure Stance where
 /-- Whether a `Stance` is not `isPublic`. -/
 @[expose, macro_inline, grind] public def Stance.isPrivate (s : Stance) := !s.isPublic
 
+def Stance.toString : Stance → String
+  | { isPublic, isExposed, isMeta, isExact, .. } =>
+    s!"{if isExact then "" else "(≥) "}\
+      {if let some isExposed := isExposed then
+        if isExposed then "@[expose] " else "@[no_expose] "
+      else ""}\
+      {if isPublic then "public" else "private"}\
+      {if let some isMeta := isMeta then
+        if isMeta then " meta" else " runtime"
+      else ""}"
+
+instance : ToString Stance := ⟨(·.toString)⟩
+
 /-- A set of declaration needs. -/
 /- TODO: make this more efficient, ideally a tiny `BitsetOf`, if such API were created. Also
 consider just a small `List` or `Array` that avoids duplication at insertion time. -/
@@ -453,7 +476,7 @@ structure DeclNeed where
   -- TODO: more efficient data structure?
   /-- Declarations whose location may not be altered. A map
   `[module name] ↦ [needs of declarations from that module]`. -/
-  fixedDecls : Std.TreeMap Name (NameMap DeclDeclNeedsKindSet) := {}
+  fixedDecls : NameMap (NameMap DeclDeclNeedsKindSet) := {}
   /--
   Extra module uses recorded in shake extensions during the command. Note that if several
   declarations were created during the command, the extra mod uses are the same among all of them.
@@ -466,6 +489,37 @@ structure DeclNeed where
   childDecls : List Name := []
 deriving Inhabited, Repr
 
+/-- Render a `NameMap` as `{ key₁ ↦ val₁, key₂ ↦ val₂, ... }`. -/
+private def NameMap.toMessageData (n : NameMap α) (f : α → MessageData)
+    (isConst := false) : MessageData :=
+  if n.isEmpty then "{}" else
+    let ns := n.toList.map fun (n, a) =>
+      m!"{if isConst then MessageData.ofConstName n else m!"{n}"} ↦ {f a}"
+    .bracket "{" ((m!"," ++ Format.line).joinSep ns) "}"
+
+def DeclNeed.toMessageData : DeclNeed → MessageData
+  | { freeDecls, fixedDecls, extraModUses, parentDecls, childDecls } => Id.run do
+    let mut msgs := #[]
+    unless freeDecls.isEmpty do
+      msgs := msgs.push m!"freeDecls := \
+        {NameMap.toMessageData freeDecls (m!"{·.toList.map (·.toMessageData)}") (isConst := true)}"
+    unless fixedDecls.isEmpty do
+      msgs := msgs.push m!"fixedDecls := \
+        {NameMap.toMessageData fixedDecls fun n =>
+          NameMap.toMessageData n (m!"{·.toList.map (·.toMessageData)}") (isConst := true)}"
+    unless extraModUses.isEmpty do
+      msgs := msgs.push m!"extraModUses := \
+        {NameMap.toMessageData extraModUses fun ks =>
+          m!"{ks.toList.map toString}"}"
+    unless parentDecls.isEmpty do
+      msgs := msgs.push m!"parentDecls := {parentDecls.map MessageData.ofConstName}"
+    unless childDecls.isEmpty do
+      msgs := msgs.push m!"parentDecls := {childDecls.map MessageData.ofConstName}"
+    if msgs.isEmpty then return "{}" else
+      return .bracket "{" ((m!"," ++ Format.line).joinSep msgs.toList) "}"
+
+instance : ToMessageData DeclNeed := ⟨(·.toMessageData)⟩
+
 /-- Records that the ambient declaration needs the declaration `decl` at availability `k`, where
 the module of `decl` is left free. -/
 @[inline] def DeclNeed.insertFreeDecl (k : DeclDeclNeedsKind) (decl : Name) (declNeed : DeclNeed) :
@@ -497,6 +551,11 @@ def DeclNeed.insertExtraModUses (extraModUses : List ExtraModUse) (declNeed : De
 -- TODO: more efficient data structure
 abbrev DeclNeeds := NameMap DeclNeed
 
+def DeclNeeds.toMessageData (needs : DeclNeeds) : MessageData :=
+  NameMap.toMessageData needs (·.toMessageData) (isConst := true)
+
+instance : ToMessageData DeclNeeds := ⟨(·.toMessageData)⟩
+
 @[inline] def DeclNeeds.insertAutoDeclLink (autoDecl parentDecl : Name) (declNeeds : DeclNeeds) :
     DeclNeeds :=
   declNeeds.alter autoDecl (fun need? =>
@@ -735,7 +794,7 @@ instead of being duplicated across different declarations. -/
 def DeclNeeds.calcSyntaxNeeds (env : Environment) (decls : Array Name) (cmd : Syntax)
     (declNeeds : DeclNeeds) : CommandElabM DeclNeeds := do
   trace[ImportGraph.Shake] "Calculating syntax needs for:\
-    {indentD <| toMessageData <| decls.map MessageData.ofConstName}"
+    {indentD m!"{decls.map MessageData.ofConstName}"}"
   let indirectModUses := indirectModUseExt.getState env
   let mut declNeeds := declNeeds
   for stx in cmd.topDown do
@@ -828,3 +887,5 @@ def withElabCommandCapturingNeeds (cmd : Syntax.Command) :
     declNeeds ← declNeeds.calcSyntaxNeeds env newDecls cmd
     declNeeds ← liftCoreM <| declNeeds.calcIRNeeds.run'
     return (declNeeds, newDecls)
+
+initialize registerTraceClass `ImportGraph.Shake

ImportGraph/Shake/Environment.lean

@@ -105,7 +105,7 @@ def toSimultaneousImportNeeds (env : Environment)
         let some usedStance ← getStance? usedDecl | continue
         let mut usedKs : DeclDeclNeedsKindSet := {}
         for k in ks do
-          trace[ImportGraph.Shake] "{k.pretty}"
+          trace[ImportGraph.Shake] "{k}"
           if usedKs.contains k then continue
           usedKs := usedKs.insert k
           if let .comptime <| .indirect _ mods := k then

ImportGraph/Shake/Workspace.lean

@@ -72,7 +72,7 @@ def DeclNeeds.toSimultaneousImportNeeds
           (fun _ => return m!"Uses decl `{.ofConstName usedDecl}`") do←
         let some usedStance ← getStance? usedDecl | continue
         for k in ks do
-          trace[ImportGraph.Shake] "{k.pretty}"
+          trace[ImportGraph.Shake] "{k}"
           if let .comptime <| .indirect _ mods := k then
             for modName in mods do
               let some modIdx := w.idxOfMod[modName]? | continue
@@ -109,5 +109,3 @@ def ImportNeeds.providersByLib (w : WorkspaceModel) (needs : ImportNeeds)
       |>.then (compare pᵢ.size pⱼ.size)
       |>.then (Name.cmp (w.getMod! i).name (w.getMod! j).name) -- for stability if all else fails
       |>.isLT).map (·.1)
-
-initialize registerTraceClass `ImportGraph.Shake

ImportGraph/Tools/FindHome.lean

@@ -147,7 +147,7 @@ elab_rules : command
       let modName := w.getMod! modIdx |>.name
       let mut decls : NameSet := {}
       for (_, need) in declNeeds do
-        let some declsFromMod := need.fixedDecls[modName]? | continue
+        let some declsFromMod := need.fixedDecls.get? modName | continue
         decls := decls.insertMany declsFromMod.keysArray
       let declsArray := decls.toArray
       links := links.push <|← goToModuleOfDecls declsArray (fallbackModule := modName)

ImportGraphTest/Shake/DeclNeeds.lean

@@ -0,0 +1,74 @@
+module
+
+public meta import ImportGraph.Shake.Workspace
+public import Lean.Elab.Command
+
+import ImportGraphTest.Shake.Algebra
+
+open ImportGraph Shake Lean Elab Command
+
+elab "#show_decl_needs" colGe cmd:command : command => do
+  let (declNeeds, newDecls) ← withElabCommandCapturingNeeds cmd
+  let w ← getWorkspaceModel (extraMods := #[← getMainModule])
+  let (importNeeds, stances) ← liftCoreM <| declNeeds.toSimultaneousImportNeeds w |>.run
+  let stanceMsg : MessageData := .bracket (l := "{") (r := "}") <|
+    (m!"," ++ Format.line).joinSep <|
+      stances.toList.mergeSort (·.1.cmp ·.1 |>.isLE) |>.map fun (n, s) =>
+        m!"{.ofConstName n} ↦ {match s with | some s => s.toString | _ => "none"}"
+  logInfo m!"\
+    New decls:{indentD (newDecls.toList.map MessageData.ofConstName)}\
+    \nNeeds:{indentD declNeeds.toMessageData}\
+    \nStances:{indentD stanceMsg}\
+    \nImportNeeds:{indentD <| importNeeds.toString}"
+
+/--
+info: New decls:
+  [x]
+Needs:
+  {x ↦ {fixedDecls := {Init.Prelude ↦ {true ↦ [used in value, used in runtime IR], Bool ↦ [used in type]}}}}
+Stances:
+  {Bool ↦ @[no_expose] public, x ↦ private runtime, true ↦ @[no_expose] public}
+ImportNeeds:
+  │⠒⠀│
+-/
+#guard_msgs in
+#show_decl_needs def x : Bool := true
+
+/--
+info: New decls:
+  [x']
+Needs:
+  {x' ↦ {fixedDecls := {Init.Prelude ↦ {true ↦ [used in value, used in runtime IR], Bool ↦ [used in type]}}}}
+Stances:
+  {Bool ↦ @[no_expose] public, x' ↦ public runtime, true ↦ @[no_expose] public}
+ImportNeeds:
+  │⠚⠀│
+-/
+#guard_msgs in
+#show_decl_needs public def x' : Bool := true
+
+/--
+info: New decls:
+  [x'']
+Needs:
+  {x'' ↦ {fixedDecls := {Init.Prelude ↦ {true ↦ [used in value, used in runtime IR], Bool ↦ [used in type]}}}}
+Stances:
+  {Bool ↦ @[no_expose] public, x'' ↦ @[expose] public runtime, true ↦ @[no_expose] public}
+ImportNeeds:
+  │⠊⠀│
+-/
+#guard_msgs in
+#show_decl_needs @[expose] public def x'' : Bool := true
+
+/--
+info: New decls:
+  [x''']
+Needs:
+  {x''' ↦ {fixedDecls := {Init.Prelude ↦ {true ↦ [used in value], Bool ↦ [used in type]}}}}
+Stances:
+  {Bool ↦ @[no_expose] public, x''' ↦ @[expose] public, true ↦ @[no_expose] public}
+ImportNeeds:
+  │⠈⠀│
+-/
+#guard_msgs in
+#show_decl_needs @[expose] public noncomputable def x''' : Bool := true