diff --git a/core/src/main/scala/org/bykn/bosatsu/Package.scala b/core/src/main/scala/org/bykn/bosatsu/Package.scala index 56d81c5ae..8abad1e4a 100644 --- 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 index 29ad31053..fcaa32a36 100644 --- 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