parser: use ast user data to store attribution

ast/src/main/kotlin/gay/pizza/pork/ast/Definition.kt

@@ -8,5 +8,6 @@ import kotlinx.serialization.Serializable
 @SerialName("definition")
 sealed class Definition : Node() {
   abstract val symbol: Symbol
+
   abstract val modifiers: DefinitionModifiers
 }

ast/src/main/kotlin/gay/pizza/pork/ast/Node.kt

@@ -1,6 +1,7 @@
 // GENERATED CODE FROM PORK AST CODEGEN
 package gay.pizza.pork.ast
 
+import kotlinx.serialization.Transient
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 
@@ -9,6 +10,9 @@ import kotlinx.serialization.Serializable
 sealed class Node {
   abstract val type: NodeType
 
+  @Transient
+  var data: Any? = null
+
   open fun <T> visitChildren(visitor: NodeVisitor<T>): List<T> =
     emptyList()
 

ast/src/main/kotlin/gay/pizza/pork/ast/NodeExtensions.kt

@@ -0,0 +1,6 @@
+// GENERATED CODE FROM PORK AST CODEGEN
+package gay.pizza.pork.ast
+
+@Suppress("UNCHECKED_CAST")
+fun <P> Node.data(): P? =
+  data as? P?

buildext/src/main/kotlin/gay/pizza/pork/buildext/ast/AstCodegen.kt

@@ -18,6 +18,7 @@ class AstCodegen(val pkg: String, val outputDirectory: Path, val world: AstWorld
     for (type in world.typeRegistry.types) {
       writeAstType(type)
     }
+    writeNodeExtensions()
     writeNodeType()
     writeNodeVisitors()
     writeNodeCoalescer()
@@ -124,7 +125,6 @@ class AstCodegen(val pkg: String, val outputDirectory: Path, val world: AstWorld
     visitAllFunction.body.add(
       "nodeLists.asSequence().flatten().filterNotNull().map { visit(it) }.toList()")
     visitorExtensionSet.functions.add(visitAllFunction)
-
     write("NodeVisitorExtensions.kt", KotlinWriter(visitorExtensionSet))
   }
 
@@ -190,12 +190,22 @@ class AstCodegen(val pkg: String, val outputDirectory: Path, val world: AstWorld
     }
 
     if (role == AstTypeRole.RootNode) {
+      kotlinClassLike.imports.add("kotlinx.serialization.Transient")
       val typeMember = KotlinMember(
         "type",
         "NodeType",
         abstract = true
       )
       kotlinClassLike.members.add(typeMember)
+      val dataMember = KotlinMember(
+        "data",
+        "Any?",
+        value = "null",
+        mutable = true,
+        notInsideConstructor = true,
+        annotations = mutableListOf("@Transient")
+      )
+      kotlinClassLike.members.add(dataMember)
 
       val abstractVisitChildrenFunction = KotlinFunction(
         "visitChildren",
@@ -389,6 +399,21 @@ class AstCodegen(val pkg: String, val outputDirectory: Path, val world: AstWorld
     write("${type.name}.kt", KotlinWriter(kotlinClassLike))
   }
 
+  private fun writeNodeExtensions() {
+    val nodeExtensionSet = KotlinFunctionSet(pkg)
+    val dataFunction = KotlinFunction(
+      "data",
+      typeParameters = mutableListOf("P"),
+      extensionOf = "Node",
+      returnType = "P?",
+      isImmediateExpression = true,
+      annotations = mutableListOf("""@Suppress("UNCHECKED_CAST")""")
+    )
+    dataFunction.body.add("data as? P?")
+    nodeExtensionSet.functions.add(dataFunction)
+    write("NodeExtensions.kt", KotlinWriter(nodeExtensionSet))
+  }
+
   private fun toKotlinType(typeRef: AstTypeRef): String {
     val baseType = typeRef.type?.name ?: typeRef.primitive?.id
       ?: throw RuntimeException("Unable to determine base type.")

buildext/src/main/kotlin/gay/pizza/pork/buildext/codegen/KotlinMember.kt

@@ -6,5 +6,9 @@ class KotlinMember(
   var abstract: Boolean = false,
   var overridden: Boolean = false,
   var value: String? = null,
-  var mutable: Boolean = false
+  var mutable: Boolean = false,
+  var private: Boolean = false,
+  var protected: Boolean = false,
+  var notInsideConstructor: Boolean = false,
+  var annotations: MutableList<String> = mutableListOf()
 )

buildext/src/main/kotlin/gay/pizza/pork/buildext/codegen/KotlinWriter.kt

@@ -15,7 +15,7 @@ class KotlinWriter() {
     }
     writeClassLike(classType, kotlinClass)
     val members = kotlinClass.members.filter {
-      it.abstract || (it.overridden && it.value != null)
+      it.abstract || (it.overridden && it.value != null) || it.notInsideConstructor
     }
     if (members.isEmpty() && kotlinClass.functions.isEmpty()) {
       appendLine()
@@ -23,13 +23,23 @@ class KotlinWriter() {
       appendLine(" {")
     }
 
-    for (member in members) {
+    for ((index, member) in members.withIndex()) {
-      val form = if (member.mutable) "var" else "val"
+      for (annotation in member.annotations) {
+        appendLine("  $annotation")
+      }
+
+      val privacy = when {
+        member.private -> "private "
+        member.protected -> "protected "
+        else -> ""
+      }
+      val form = if (member.mutable) "${privacy}var" else "${privacy}val"
       if (member.abstract) {
         appendLine("  abstract $form ${member.name}: ${member.type}")
       } else {
+        append("  ")
         if (member.overridden) {
-          append("  override ")
+          append("override ")
         }
         append("$form ${member.name}: ${member.type}")
         if (member.value != null) {
@@ -38,6 +48,10 @@ class KotlinWriter() {
         }
         appendLine()
       }
+
+      if (index != members.size - 1) {
+        appendLine()
+      }
     }
 
     if (members.isNotEmpty() && kotlinClass.functions.isNotEmpty()) {
@@ -109,7 +123,7 @@ class KotlinWriter() {
     }
 
     val contructedMembers = kotlinClass.members.filter {
-      !it.abstract && !(it.overridden && it.value != null)
+      !it.abstract && !(it.overridden && it.value != null) && !it.notInsideConstructor
     }
 
     if (contructedMembers.isNotEmpty()) {

parser/src/main/kotlin/gay/pizza/pork/parser/Parser.kt

@@ -2,6 +2,7 @@ package gay.pizza.pork.parser
 
 import gay.pizza.pork.ast.*
 
+@Suppress("SameParameterValue")
 class Parser(source: PeekableSource<Token>, val attribution: NodeAttribution) {
   private val unsanitizedSource = source
 

parser/src/main/kotlin/gay/pizza/pork/parser/ParserAttributes.kt

@@ -0,0 +1,23 @@
+package gay.pizza.pork.parser
+
+import gay.pizza.pork.ast.Node
+import gay.pizza.pork.ast.NodeCoalescer
+import gay.pizza.pork.ast.data
+import gay.pizza.pork.ast.visit
+
+data class ParserAttributes(val tokens: List<Token>) {
+  companion object {
+    fun recallAllTokens(node: Node): List<Token> {
+      val all = mutableListOf<Token>()
+      val coalescer = NodeCoalescer { item ->
+        val attributes = item.data<ParserAttributes>()
+        if (attributes != null) {
+          all.addAll(attributes.tokens)
+        }
+      }
+      coalescer.visit(node)
+      all.sortBy { it.start }
+      return all
+    }
+  }
+}

parser/src/main/kotlin/gay/pizza/pork/parser/ParserNodeAttribution.kt

@@ -0,0 +1,37 @@
+package gay.pizza.pork.parser
+
+import gay.pizza.pork.ast.Node
+import gay.pizza.pork.ast.data
+
+class ParserNodeAttribution : NodeAttribution {
+  private val stack = mutableListOf<MutableList<Token>>()
+  private var current: MutableList<Token>? = null
+
+  override fun enter() {
+    val store = mutableListOf<Token>()
+    current = store
+    stack.add(store)
+  }
+
+  override fun push(token: Token) {
+    val store = current ?: throw RuntimeException("enter() not called!")
+    store.add(token)
+  }
+
+  override fun <T : Node> adopt(node: T) {
+    val attributes = node.data<ParserAttributes>()
+    if (attributes != null) {
+      for (token in attributes.tokens) {
+        push(token)
+      }
+      node.data = ParserAttributes(emptyList())
+    }
+  }
+
+  override fun <T: Node> exit(node: T): T {
+    val store = stack.removeLast()
+    current = stack.lastOrNull()
+    node.data = ParserAttributes(store)
+    return node
+  }
+}

parser/src/main/kotlin/gay/pizza/pork/parser/TokenNodeAttribution.kt

@@ -1,54 +0,0 @@
-package gay.pizza.pork.parser
-
-import gay.pizza.pork.ast.NodeCoalescer
-import gay.pizza.pork.ast.Node
-import gay.pizza.pork.ast.visit
-import java.util.IdentityHashMap
-
-class TokenNodeAttribution : NodeAttribution {
-  val nodes: MutableMap<Node, List<Token>> = IdentityHashMap()
-
-  private val stack = mutableListOf<MutableList<Token>>()
-  private var current: MutableList<Token>? = null
-
-  override fun enter() {
-    val store = mutableListOf<Token>()
-    current = store
-    stack.add(store)
-  }
-
-  override fun push(token: Token) {
-    val store = current ?: throw RuntimeException("enter() not called!")
-    store.add(token)
-  }
-
-  override fun <T : Node> adopt(node: T) {
-    val tokens = nodes.remove(node)
-    if (tokens != null) {
-      for (token in tokens) {
-        push(token)
-      }
-    }
-  }
-
-  override fun <T: Node> exit(node: T): T {
-    val store = stack.removeLast()
-    nodes[node] = store
-    current = stack.lastOrNull()
-    return node
-  }
-
-  fun tokensOf(node: Node): List<Token>? = nodes[node]
-
-  fun assembleTokens(node: Node): List<Token> {
-    val allTokens = mutableListOf<Token>()
-    val coalescer = NodeCoalescer { item ->
-      val tokens = tokensOf(item)
-      if (tokens != null) {
-        allTokens.addAll(tokens)
-      }
-    }
-    coalescer.visit(node)
-    return allTokens.asSequence().distinct().sortedBy { it.start }.toList()
-  }
-}

tool/src/main/kotlin/gay/pizza/pork/tool/AttributeCommand.kt

@@ -3,23 +3,26 @@ package gay.pizza.pork.tool
 import com.github.ajalt.clikt.core.CliktCommand
 import com.github.ajalt.clikt.parameters.arguments.argument
 import gay.pizza.dough.fs.PlatformFsProvider
+import gay.pizza.pork.ast.Node
 import gay.pizza.pork.ast.NodeCoalescer
+import gay.pizza.pork.ast.data
 import gay.pizza.pork.ast.visit
 import gay.pizza.pork.minimal.FileTool
-import gay.pizza.pork.parser.TokenNodeAttribution
+import gay.pizza.pork.parser.ParserAttributes
+import gay.pizza.pork.parser.ParserNodeAttribution
 
 class AttributeCommand : CliktCommand(help = "Attribute AST", name = "attribute") {
   val path by argument("file")
 
   override fun run() {
     val tool = FileTool(PlatformFsProvider.resolve(path))
-    val attribution = TokenNodeAttribution()
+    val attribution = ParserNodeAttribution()
     val compilationUnit = tool.parse(attribution)
 
     val coalescer = NodeCoalescer { node ->
-      val tokens = attribution.assembleTokens(node)
+      val allTokens = ParserAttributes.recallAllTokens(node)
       println("node ${node.type.name}")
-      for (token in tokens) {
+      for (token in allTokens) {
         println("token $token")
       }
     }