Add methods to remove unused values (treeshake) (#1273)

* Add methods to remove unused values (treeshake) * minor cleanup
johnynek · Nov 22, 2024 · a4bbaa6 · a4bbaa6
1 parent 09e3e62
commit a4bbaa6
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 5 deletions.
diff --git a/core/src/main/scala/org/bykn/bosatsu/Package.scala b/core/src/main/scala/org/bykn/bosatsu/Package.scala
@@ -64,16 +64,17 @@ object Package {
   type Parsed = Package[PackageName, Unit, Unit, List[Statement]]
   type Resolved =
     FixPackage[Unit, Unit, (List[Statement], ImportMap[PackageName, Unit])]
+  type TypedProgram[T] = (
+        Program[TypeEnv[Kind.Arg], TypedExpr[T], Any],
+        ImportMap[Interface, NonEmptyList[Referant[Kind.Arg]]]
+    )
   type Typed[T] = Package[
     Interface,
     NonEmptyList[Referant[Kind.Arg]],
     Referant[
       Kind.Arg
     ],
-    (
-        Program[TypeEnv[Kind.Arg], TypedExpr[T], Any],
-        ImportMap[Interface, NonEmptyList[Referant[Kind.Arg]]]
-    )
+    TypedProgram[T]
   ]
   type Inferred = Typed[Declaration]
 
@@ -528,6 +529,12 @@ object Package {
       (Package.Interface, ImportedName[NonEmptyList[Referant[Kind.Arg]]])
     ] =
       pack.program._2(n)
+
+    def filterLets(fn: Identifier => Boolean): Typed[A] = {
+      val (prog, importMap) = pack.program
+      val prog1 = prog.copy(lets = prog.lets.filter { case (b, _, _) => fn(b) })
+      pack.copy(program = (prog1, importMap))
+    }
   }
 
   def orderByName[A, B, C, D]: Order[Package[A, B, C, D]] =

diff --git a/core/src/main/scala/org/bykn/bosatsu/PackageMap.scala b/core/src/main/scala/org/bykn/bosatsu/PackageMap.scala
@@ -1,6 +1,6 @@
 package org.bykn.bosatsu
 
-import org.bykn.bosatsu.graph.{Memoize, Toposort}
+import org.bykn.bosatsu.graph.{Dag, Memoize, Toposort}
 import cats.{Foldable, Monad, Show}
 import cats.data.{
   Ior,
@@ -111,6 +111,26 @@ object PackageMap {
   // convenience for type inference
   def toAnyTyped[A](p: Typed[A]): Typed[Any] = p
 
+  def treeShake[A](p: Typed[A], roots: Set[(PackageName, Identifier)]): Typed[A] = {
+    type Ident = (PackageName, Identifier)
+
+    def dependency(a: Ident): Iterable[Ident] =
+      p.toMap.get(a._1) match {
+        case Some(pack) =>
+          pack.lets.flatMap { case (_, _, te) => te.globals }
+        case None => Nil
+      }
+
+    val keep = Dag.transitiveSet(roots.toList.sorted)(dependency)
+
+    val kept = p.toMap.iterator.map { case (_, pack) =>
+      pack.filterLets { nm => keep((pack.name, nm)) }
+    }
+    .toList
+
+    fromIterable(kept)
+  }
+
   type Inferred = Typed[Declaration]
 
   /** This builds a DAG of actual packages where names have been replaced by the