feat(dataflow): sparse fact api

Veir/Analysis/DataFlow/Facts.lean

@@ -82,6 +82,12 @@ Stored in the entry block of each region.
 structure RegionMetadataPayload where
   postOrderIndex : HashMap BlockPtr Nat := {}
 
+/--
+A sparse dataflow fact payload for one abstract domain.
+-/
+structure SparsePayload (Domain : Type) where
+  latticeElement : Domain
+
 /--
 The fact specific data stored for each fact kind.
 -/
@@ -92,12 +98,15 @@ The fact specific data stored for each fact kind.
 /--
 A dataflow fact stored by the framework.
 
-Each fact associates with a lattice anchor (some location in the program), has 
+Each fact associates with a lattice anchor (some location in the program), has
 an array of dependents (other facts that "depend" on this fact's current state in
-some fashion), and has the fact specific payload determined by its `FactKind`.
+some fashion), has an array of analysis subscribers (similar to dependents except 
+it's entire analyses that depend on this fact's current state), and has the fact 
+specific payload determined by its `FactKind`.
 -/
 structure Fact (kind : FactKind) where
   dependents : Array WorkItem := #[]
+  subscribers : Array AnalysisKind := #[]
   payload : FactPayload kind
 
 namespace Fact
@@ -114,6 +123,15 @@ Add one dependent work item to the fact.
 def addDependent (fact : Fact kind) (workItem : WorkItem) : Fact kind :=
   fact.setDependents (fact.dependents.push workItem)
 
+/--
+Subscribe one analysis to changes of this fact.
+-/
+def subscribe (fact : Fact kind) (analysisKind : AnalysisKind) : Fact kind :=
+  if fact.subscribers.contains analysisKind then
+    fact
+  else
+    { fact with subscribers := (fact.subscribers.push analysisKind) }
+
 /--
 Enqueue all dependents of this fact.
 -/

Veir/Analysis/DataFlow/SparseFact.lean

@@ -0,0 +1,83 @@
+module
+
+public import Veir.Analysis.DataFlowFramework
+public import Veir.Analysis.DataFlow.Domains.AbstractDomain
+
+public section
+
+namespace Veir
+
+/--
+Implement this class to register a custom type to be recognized as
+a sparse fact type by the dataflow framework.
+-/
+class SparseFactSpec (kind : FactKind) (Domain : outParam Type) where
+  payloadEq : FactPayload kind = SparsePayload Domain
+
+namespace SparseFact
+
+variable {kind : FactKind} {Domain : Type}
+variable [SparseFactSpec kind Domain]
+
+def getPayload (fact : Fact kind) : SparsePayload Domain :=
+  cast SparseFactSpec.payloadEq fact.payload
+
+def setPayload (fact : Fact kind) (payload : SparsePayload Domain) : Fact kind :=
+  { fact with payload := cast (Eq.symm SparseFactSpec.payloadEq) payload }
+
+def latticeElement (fact : Fact kind) : Domain :=
+  (getPayload fact).latticeElement
+
+def setLatticeElement (fact : Fact kind) (latticeElement : Domain) : Fact kind :=
+  let payload := getPayload fact
+  setPayload fact { payload with latticeElement := latticeElement }
+
+/--
+Propagate a sparse lattice update by revisiting dependents and all users of the
+updated SSA value for subscribed analyses.
+-/
+def propagate (state : Fact kind) (anchor : LatticeAnchor) 
+  (dfCtx : DataFlowContext) (irCtx : IRContext OpCode) : DataFlowContext := Id.run do
+  let mut dfCtx := { dfCtx with workList := state.enqueueDependents dfCtx.workList }
+  match anchor with
+  | .ValuePtr ssaValue =>
+    let mut maybeUse := ssaValue.getFirstUse! irCtx
+    while let some use := maybeUse do
+      let user := (use.get! irCtx).owner
+      match InsertPoint.after? user irCtx with
+      | some point =>
+        for analysisKind in state.subscribers do
+          dfCtx := dfCtx.enqueue (point, analysisKind)
+      | none => pure ()
+      maybeUse := (use.get! irCtx).nextUse
+  | _ =>
+    pure ()
+  dfCtx
+
+section
+
+variable [Bot Domain]
+
+/-- Default sparse lattice fact for the given anchor. -/
+def mkDefault : Fact kind :=
+  { payload := cast (Eq.symm SparseFactSpec.payloadEq) { latticeElement := ⊥ } }
+
+instance : FactSpec kind where
+  mkDefault := SparseFact.mkDefault (kind := kind)
+  propagate := SparseFact.propagate (kind := kind)
+
+end
+
+def getElement? (kind : FactKind) [SparseFactSpec kind Domain] [FactSpec kind]
+    (ssaValue : ValuePtr) (dfCtx : DataFlowContext) : Option Domain := do
+  let state ← dfCtx.getFact? kind (.ValuePtr ssaValue)
+  return latticeElement state
+
+def getElementD (kind : FactKind) [SparseFactSpec kind Domain] [FactSpec kind]
+    (ssaValue : ValuePtr) (fallback : Domain)
+    (dfCtx : DataFlowContext) : Domain :=
+  (getElement? kind ssaValue dfCtx).getD fallback
+
+end SparseFact
+
+end Veir