ir gen mvp

hail/build.mill

@@ -73,7 +73,7 @@ object Deps {
   }
 }
 
-trait HailScalaModule extends SbtModule with ScalafmtModule with ScalafixModule { outer =>
+trait HailScalaModule extends ScalafmtModule with ScalafixModule { outer =>
   override def scalaVersion: T[String] = build.env.scalaVersion()
 
   override def javacOptions: T[Seq[String]] = Seq(
@@ -107,7 +107,7 @@ trait HailScalaModule extends SbtModule with ScalafmtModule with ScalafixModule
   override def bspCompileClasspath: T[Agg[UnresolvedPath]] =
     super.bspCompileClasspath() ++ resources().map(p => UnresolvedPath.ResolvedPath(p.path))
 
-  trait HailTests extends SbtTests with TestNg with ScalafmtModule with ScalafixModule {
+  trait HailTests extends ScalaTests with TestNg with ScalafmtModule with ScalafixModule {
     override def forkArgs: T[Seq[String]] = Seq("-Xss4m", "-Xmx4096M")
 
     override def ivyDeps: T[Agg[Dep]] =
@@ -125,7 +125,7 @@ trait HailScalaModule extends SbtModule with ScalafmtModule with ScalafixModule
   }
 }
 
-object `package` extends RootModule with HailScalaModule { outer =>
+object `package` extends RootModule with HailScalaModule with SbtModule { outer =>
 
   override def millSourcePath: Path = super.millSourcePath / "modules"
 
@@ -179,6 +179,10 @@ object `package` extends RootModule with HailScalaModule { outer =>
     buildInfo(),
   )
 
+  override def generatedSources: T[Seq[PathRef]] = Task {
+    Seq(`ir-gen`.generate())
+  }
+
   override def unmanagedClasspath: T[Agg[PathRef]] =
     Agg(shadedazure.assembly())
 
@@ -250,6 +254,19 @@ object `package` extends RootModule with HailScalaModule { outer =>
     PathRef(T.dest)
   }
 
+  object `ir-gen` extends HailScalaModule {
+    def ivyDeps = Agg(
+      ivy"com.lihaoyi::mainargs:0.6.2",
+      ivy"com.lihaoyi::os-lib:0.10.7",
+      ivy"com.lihaoyi::sourcecode:0.4.2",
+    )
+
+    def generate: T[PathRef] = Task {
+      runForkedTask(finalMainClass, Task.Anon { Args("--path", T.dest) })()
+      PathRef(T.dest)
+    }
+  }
+
   object memory extends JavaModule { // with CrossValue {
     override def zincIncrementalCompilation: T[Boolean] = false
 
@@ -267,7 +284,7 @@ object `package` extends RootModule with HailScalaModule { outer =>
     }
   }
 
-  object test extends HailTests {
+  object test extends HailTests with SbtTests {
     override def resources: T[Seq[PathRef]] = outer.resources() ++ super.resources()
 
     override def assemblyRules: Seq[Rule] = outer.assemblyRules ++ Seq(

hail/modules/ir-gen/src/Main.scala

@@ -0,0 +1,51 @@
+import mainargs.{ParserForMethods, main}
+
+import scala.collection.mutable
+
+object NodeRegistry {
+  val allNodes: mutable.Set[IR] = mutable.Set.empty[IR]
+}
+
+sealed abstract class AttOrChild {
+  val name: String
+  def generateDeclaration: String
+}
+
+final case class Att(name: String, typ: String) extends AttOrChild {
+  override def generateDeclaration: String = s"$name: $typ"
+}
+
+final case class Child(name: String) extends AttOrChild {
+  override def generateDeclaration: String = s"$name: IR"
+}
+
+sealed abstract class IR(implicit val name: sourcecode.Name) {
+  NodeRegistry.allNodes += this
+
+  val attsAndChildren: Seq[AttOrChild]
+  val isTrivial: Boolean = false
+
+  def generateDef: String =
+    (s"final case class $name(${attsAndChildren.map(_.generateDeclaration).mkString(",")}) extends IR"
+      + (if (isTrivial) " with TrivialIR" else ""))
+}
+
+object I32 extends IR {
+  override val attsAndChildren = Seq(Att("x", "Int"))
+  override val isTrivial = true
+}
+
+object I64 extends IR {
+  override val attsAndChildren = Seq(Att("x", "Int"))
+  override val isTrivial = true
+}
+
+object Main {
+  @main
+  def main(path: String) = {
+    val gen = "package is.hail.expr.ir\n" + NodeRegistry.allNodes.map(_.generateDef).mkString("\n")
+    os.write(os.Path(path) / "IR_gen.scala", gen)
+  }
+
+  def main(args: Array[String]): Unit = ParserForMethods(this).runOrExit(args)
+}

hail/modules/src/main/scala/is/hail/expr/ir/IR.scala

@@ -27,7 +27,7 @@ import java.io.OutputStream
 import org.json4s.{DefaultFormats, Extraction, Formats, JValue, ShortTypeHints}
 import org.json4s.JsonAST.{JNothing, JString}
 
-sealed trait IR extends BaseIR {
+trait IR extends BaseIR {
   private var _typ: Type = null
 
   def typ: Type = {
@@ -72,12 +72,12 @@ sealed trait IR extends BaseIR {
   def unwrap: IR = _unwrap(this)
 }
 
-sealed trait TypedIR[T <: Type] extends IR {
+trait TypedIR[T <: Type] extends IR {
   override def typ: T = tcoerce[T](super.typ)
 }
 
 // Mark Refs and constants as IRs that are safe to duplicate
-sealed trait TrivialIR extends IR
+trait TrivialIR extends IR
 
 object Literal {
   def coerce(t: Type, x: Any): IR = {
@@ -146,7 +146,7 @@ class WrappedByteArrays(val ba: Array[Array[Byte]]) {
   }
 }
 
-final case class I32(x: Int) extends IR with TrivialIR
+//final case class I32(x: Int) extends IR with TrivialIR
 final case class I64(x: Long) extends IR with TrivialIR
 final case class F32(x: Float) extends IR with TrivialIR
 final case class F64(x: Double) extends IR with TrivialIR