Some cleanups and refactors to visitor trait.

src/main/scala/lexi/Tree.scala

@@ -4,10 +4,6 @@ trait Tree {
   def parent: Option[Tree]
   def children: List[Tree]
 
-  def accept(visitor: Visitor) =
+  def accept(visitor: Visitor[Tree]) =
     visitor.visit(this)
 }
-
-trait Visitor {
-  def visit(tree: Tree): Tree
-}

src/main/scala/lexi/Visitor.scala

@@ -0,0 +1,5 @@
+package lexi
+
+trait Visitor[T] {
+  def visit(tree: Tree): T
+}

src/main/scala/lexi/backends/asm/ir/IrClassVisitor.scala

@@ -1,30 +1,27 @@
 package lexi.backends.asm.ir
 
-import lexi.ir.{IrClass, IrTree}
+import lexi.{Tree, Visitor}
+import lexi.ir.IrClass
 import org.objectweb.asm.{ClassWriter, Opcodes}
 
-object IrClassVisitor extends IrVisitor[ClassWriter] {
-  override def visit(ir: IrTree): ClassWriter =
-    (((_: IrClass) => new ClassWriter(ClassWriter.COMPUTE_MAXS))
-      andThen ((writer) => {
-        writer.visit(
-          Opcodes.V1_8,
-          Opcodes.ACC_PUBLIC,
-          ir.asInstanceOf[IrClass].name.get,
-          null,
-          "java/lang/Object",
-          null
-        )
-        ir.asInstanceOf[IrClass]
-          .classBody
-          .get
-          .declarations
-          .get
-          .map(_.functionDeclaration)
-          .filter(_ != None)
-          .map(_.get)
-          .foreach(IrMethodVisitor(writer).visit(_))
-        writer.visitEnd()
-        writer
-      }))(ir.asInstanceOf[IrClass])
+object IrClassVisitor extends Visitor[ClassWriter] {
+  override def visit(ir: Tree): ClassWriter = {
+    val irClass = ir.asInstanceOf[IrClass]
+    new ClassWriter(ClassWriter.COMPUTE_MAXS) {
+      this.visit(
+        Opcodes.V1_8,
+        Opcodes.ACC_PUBLIC,
+        ir.asInstanceOf[IrClass].name.get,
+        null,
+        "java/lang/Object",
+        null
+      )
+      irClass.classBody.get.declarations.get
+        .map(_.functionDeclaration)
+        .filter(_ != None)
+        .map(_.get)
+        .foreach(IrMethodVisitor(this).visit(_))
+      visitEnd()
+    }
+  }
 }

src/main/scala/lexi/backends/asm/ir/IrMethodVisitor.scala

@@ -1,43 +1,42 @@
 package lexi.backends.asm.ir
 
-import lexi.ir.{IrFunction, IrTree}
+import lexi.{Tree, Visitor}
+import lexi.ir.IrFunction
 import org.objectweb.asm.{ClassWriter, MethodVisitor, Opcodes}
 
 object IrMethodVisitor {
   def apply(classWriter: ClassWriter): IrMethodVisitor =
     new IrMethodVisitor(classWriter)
 }
 
-class IrMethodVisitor(classWriter: ClassWriter) extends IrVisitor[MethodVisitor] {
-  override def visit(ir: IrTree): MethodVisitor =
-    (
-      (function: IrFunction) => {
-        val method = classWriter.visitMethod(
-          Opcodes.ACC_PUBLIC + Opcodes.ACC_STATIC,
-          function.name.get,
-          s"([Ljava/lang/${function.`type`};)V",
-          null,
-          null
-        )
-        method.visitCode()
-        method.visitFieldInsn(
-          Opcodes.GETSTATIC,
-          "java/lang/System",
-          "out",
-          "Ljava/io/PrintStream;"
-        )
-        method.visitLdcInsn("Class compiled!")
-        method.visitMethodInsn(
-          Opcodes.INVOKEVIRTUAL,
-          "java/io/PrintStream",
-          "println",
-          "(Ljava/lang/String;)V",
-          false
-        )
-        method.visitInsn(Opcodes.RETURN)
-        method.visitMaxs(0, 0)
-        method.visitEnd()
-        method
-      }
-    )(ir.asInstanceOf[IrFunction])
+class IrMethodVisitor(classWriter: ClassWriter) extends Visitor[MethodVisitor] {
+  override def visit(ir: Tree): MethodVisitor = {
+    val irFunction = ir.asInstanceOf[IrFunction]
+    val method = classWriter.visitMethod(
+      Opcodes.ACC_PUBLIC + Opcodes.ACC_STATIC,
+      irFunction.name.get,
+      s"([Ljava/lang/${irFunction.`type`};)V",
+      null,
+      null
+    )
+    method.visitCode()
+    method.visitFieldInsn(
+      Opcodes.GETSTATIC,
+      "java/lang/System",
+      "out",
+      "Ljava/io/PrintStream;"
+    )
+    method.visitLdcInsn("Class compiled!")
+    method.visitMethodInsn(
+      Opcodes.INVOKEVIRTUAL,
+      "java/io/PrintStream",
+      "println",
+      "(Ljava/lang/String;)V",
+      false
+    )
+    method.visitInsn(Opcodes.RETURN)
+    method.visitMaxs(0, 0)
+    method.visitEnd()
+    method
+  }
 }

src/main/scala/lexi/frontends/kotlin/phases/Ir.scala

@@ -1,18 +1,20 @@
 package lexi.frontends.kotlin.phases
 
 import lexi.{Phase, Tree, Visitor}
-import lexi.frontends.kotlin.transformations.ir.{KtClassVisitor, KtFileVisitor, KtFunctionVisitor, KtPropertyVisitor}
+import lexi.frontends.kotlin.transformations.ir.{
+  KtClassVisitor,
+  KtFileVisitor,
+  KtFunctionVisitor,
+  KtPropertyVisitor
+}
 import lexi.frontends.kotlin.{AST, KtClass, KtFile, KtFunction, KtProperty}
 import lexi.ir.IrTree
 
 object Ir extends Phase {
-  def apply(tree: Tree): Tree =
-    visitorType(tree).visit(tree)
-
-  private def visitorType(tree: Tree): Visitor = tree match {
+  def apply(tree: Tree): Tree = tree.match {
     case _: KtFile     => KtFileVisitor
     case _: KtClass    => KtClassVisitor
     case _: KtFunction => KtFunctionVisitor
     case _: KtProperty => KtPropertyVisitor
-  }
+  }.visit(tree)
 }

src/main/scala/lexi/frontends/kotlin/transformations/ir/KtBlockVisitor.scala

@@ -1,11 +1,12 @@
 package lexi.frontends.kotlin.transformations.ir
 
-import lexi.{Tree, Visitor}
 import lexi.frontends.kotlin.KtBlock
+import lexi.ir.IrBlock
+import lexi.{Tree, Visitor}
 
-object KtBlockVisitor extends Visitor {
-  override def visit(ast: Tree): Tree = {
+object KtBlockVisitor extends Visitor[IrBlock] {
+  override def visit(ast: Tree): IrBlock = {
     val ktBlock = ast.asInstanceOf[KtBlock]
-    new lexi.ir.IrTree {}
+    new lexi.ir.IrBlock {}
   }
 }

src/main/scala/lexi/frontends/kotlin/transformations/ir/KtClassBodyVisitor.scala

@@ -6,7 +6,7 @@ import lexi.ir.IrClassBody
 
 import scala.util.Try
 
-object KtClassBodyVisitor extends Visitor {
+object KtClassBodyVisitor extends Visitor[IrClassBody] {
   override def visit(ast: Tree): IrClassBody = {
     val ktClassBody = ast.asInstanceOf[KtClassBody]
     new IrClassBody(

src/main/scala/lexi/frontends/kotlin/transformations/ir/KtClassVisitor.scala

@@ -6,7 +6,7 @@ import lexi.ir.IrClass
 
 import scala.util.Try
 
-object KtClassVisitor extends Visitor {
+object KtClassVisitor extends Visitor[IrClass] {
   override def visit(ast: Tree): IrClass = {
     val ktClass = ast.asInstanceOf[KtClass]
     new IrClass(

src/main/scala/lexi/frontends/kotlin/transformations/ir/KtDeclarationVisitor.scala

@@ -4,7 +4,7 @@ import lexi.{Tree, Visitor}
 import lexi.frontends.kotlin.KtDeclaration
 import lexi.ir.IrDeclaration
 
-object KtDeclarationVisitor extends Visitor {
+object KtDeclarationVisitor extends Visitor[IrDeclaration] {
   override def visit(ast: Tree): IrDeclaration = {
     val ktDeclaration = ast.asInstanceOf[KtDeclaration]
     new IrDeclaration(

src/main/scala/lexi/frontends/kotlin/transformations/ir/KtExpressionVisitor.scala

@@ -4,7 +4,7 @@ import lexi.{Tree, Visitor}
 import lexi.frontends.kotlin.KtExpressionContext
 import lexi.ir.IrExpression
 
-object KtExpressionVisitor extends Visitor {
+object KtExpressionVisitor extends Visitor[IrExpression] {
   override def visit(ast: Tree): IrExpression = {
     val exprCtx = ast.asInstanceOf[KtExpressionContext]
     IrExpression(

src/main/scala/lexi/frontends/kotlin/transformations/ir/KtFileVisitor.scala

@@ -6,7 +6,7 @@ import lexi.ir.IrFile
 
 import scala.util.Try
 
-object KtFileVisitor extends Visitor {
+object KtFileVisitor extends Visitor[IrFile] {
   override def visit(tree: Tree): IrFile = {
     val ktFile = tree.asInstanceOf[KtFile]
     new IrFile(

src/main/scala/lexi/frontends/kotlin/transformations/ir/KtFunctionBodyVisitor.scala

@@ -6,7 +6,7 @@ import lexi.ir.IrFunctionBody
 
 import scala.util.Try
 
-object KtFunctionBodyVisitor extends Visitor {
+object KtFunctionBodyVisitor extends Visitor[IrFunctionBody] {
   override def visit(ast: Tree): IrFunctionBody = {
     val body = ast.asInstanceOf[KtFunctionBody]
     IrFunctionBody(

src/main/scala/lexi/frontends/kotlin/transformations/ir/KtFunctionVisitor.scala

@@ -4,7 +4,7 @@ import lexi.{Tree, Visitor}
 import lexi.frontends.kotlin.KtFunction
 import lexi.ir.IrFunction
 
-object KtFunctionVisitor extends Visitor {
+object KtFunctionVisitor extends Visitor[IrFunction] {
   override def visit(ast: Tree): IrFunction = {
     val function = ast.asInstanceOf[KtFunction]
     IrFunction(

src/main/scala/lexi/frontends/kotlin/transformations/ir/KtPropertyVisitor.scala

@@ -4,7 +4,7 @@ import lexi.{Tree, Visitor}
 import lexi.frontends.kotlin.KtProperty
 import lexi.ir.IrProperty
 
-object KtPropertyVisitor extends Visitor {
+object KtPropertyVisitor extends Visitor[IrProperty] {
   override def visit(ast: Tree): IrProperty = {
     val property = ast.asInstanceOf[KtProperty]
     IrProperty(

src/main/scala/lexi/frontends/kotlin/transformations/ir/KtTopLevelObjectVisitor.scala

@@ -4,7 +4,7 @@ import lexi.{Tree, Visitor}
 import lexi.frontends.kotlin.KtTopLevelObject
 import lexi.ir.IrTopLevelObject
 
-object KtTopLevelObjectVisitor extends Visitor {
+object KtTopLevelObjectVisitor extends Visitor[IrTopLevelObject] {
   override def visit(ast: Tree): IrTopLevelObject = {
     val tlo = ast.asInstanceOf[KtTopLevelObject]
     IrTopLevelObject(

src/main/scala/lexi/ir/Trees.scala

@@ -48,6 +48,14 @@ case class IrFunctionBody(
   expression: Option[IrExpression] = None
 ) extends IrTree
 
+case class IrBlock(
+  var statements: Option[Vector[IrStatement]] = None
+) extends IrTree
+
+case class IrStatement(
+  var expression: Option[IrExpression] = None
+) extends IrTree
+
 case class IrExpression() extends IrTree
 
 case class IrCall(