From 8647a8290ec1e40b1d180e88bcc05df12138b5ad Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Fri, 2 Oct 2020 22:21:18 +0200
Subject: [PATCH] fix code generation for using struct vars in arrays and such

---
 .../prog8/ast/expressions/AstExpressions.kt   | 12 ++++++
 .../compiler/target/c64/codegen/AsmGen.kt     |  4 +-
 .../codegen/assignment/AssignmentAsmGen.kt    | 13 +------
 examples/test.p8                              | 38 ++++++++++---------
 4 files changed, 35 insertions(+), 32 deletions(-)

diff --git a/compiler/src/prog8/ast/expressions/AstExpressions.kt b/compiler/src/prog8/ast/expressions/AstExpressions.kt
index eaeb0c973..a94697e6b 100644
--- a/compiler/src/prog8/ast/expressions/AstExpressions.kt
+++ b/compiler/src/prog8/ast/expressions/AstExpressions.kt
@@ -774,6 +774,18 @@ data class IdentifierReference(val nameInSource: List<String>, override val posi
             else -> throw FatalAstException("requires a reference value")
         }
     }
+
+    fun firstStructVarName(namespace: INameScope): String? {
+        // take the name of the first struct member of the structvariable instead
+        // if it's just a regular variable, return null.
+        val struct = memberOfStruct(namespace) ?: return null
+        val decl = targetVarDecl(namespace)!!
+        val firstStructMember = struct.nameOfFirstMember()
+        // find the flattened var that belongs to this first struct member
+        val firstVarName = listOf(decl.name, firstStructMember)
+        val firstVar = definingScope().lookup(firstVarName, this) as VarDecl
+        return firstVar.name
+    }
 }
 
 class FunctionCall(override var target: IdentifierReference,
diff --git a/compiler/src/prog8/compiler/target/c64/codegen/AsmGen.kt b/compiler/src/prog8/compiler/target/c64/codegen/AsmGen.kt
index 53e6794cb..f23986eed 100644
--- a/compiler/src/prog8/compiler/target/c64/codegen/AsmGen.kt
+++ b/compiler/src/prog8/compiler/target/c64/codegen/AsmGen.kt
@@ -411,10 +411,10 @@ internal class AsmGen(private val program: Program,
                         "$" + it.number.toInt().toString(16).padStart(4, '0')
                     }
                     is AddressOf -> {
-                        it.identifier.nameInSource.joinToString(".")
+                        it.identifier.firstStructVarName(program.namespace) ?: asmSymbolName(it.identifier)
                     }
                     is IdentifierReference -> {
-                        it.nameInSource.joinToString(".")
+                        it.firstStructVarName(program.namespace) ?: asmSymbolName(it)
                     }
                     else -> throw AssemblyError("weird array elt dt")
                 }
diff --git a/compiler/src/prog8/compiler/target/c64/codegen/assignment/AssignmentAsmGen.kt b/compiler/src/prog8/compiler/target/c64/codegen/assignment/AssignmentAsmGen.kt
index b8b9a9de8..a10d2e81c 100644
--- a/compiler/src/prog8/compiler/target/c64/codegen/assignment/AssignmentAsmGen.kt
+++ b/compiler/src/prog8/compiler/target/c64/codegen/assignment/AssignmentAsmGen.kt
@@ -320,18 +320,7 @@ internal class AssignmentAsmGen(private val program: Program, private val asmgen
     }
 
     private fun assignAddressOf(target: AsmAssignTarget, name: IdentifierReference) {
-        val struct = name.memberOfStruct(program.namespace)
-        val sourceName = if (struct != null) {
-            // take the address of the first struct member instead
-            val decl = name.targetVarDecl(program.namespace)!!
-            val firstStructMember = struct.nameOfFirstMember()
-            // find the flattened var that belongs to this first struct member
-            val firstVarName = listOf(decl.name, firstStructMember)
-            val firstVar = name.definingScope().lookup(firstVarName, name) as VarDecl
-            firstVar.name
-        } else {
-            asmgen.fixNameSymbols(name.nameInSource.joinToString("."))
-        }
+        val sourceName = name.firstStructVarName(program.namespace) ?: asmgen.fixNameSymbols(name.nameInSource.joinToString("."))
 
         when(target.kind) {
             TargetStorageKind.VARIABLE -> {
diff --git a/examples/test.p8 b/examples/test.p8
index 07ec463c0..e0d02874e 100644
--- a/examples/test.p8
+++ b/examples/test.p8
@@ -6,34 +6,36 @@
 
 main {
 
-        struct Color {
-            ubyte red
-            ubyte green
-            ubyte blue
-        }
-
+    struct Color {
+        ubyte red
+        ubyte green
+        ubyte blue
+    }
 
     Color c1 = [11,22,33]
     Color c2 = [11,22,33]
     Color c3 = [11,22,33]
-    uword[] colors = [ c1, c2, c3]      ; TODO should contain pointers to (the first element) of each struct
+    uword[] colors = [ c1, c2, c3]
 
 
     sub start() {
 
-        ubyte a
-        ubyte b
-        ubyte c
-
-        c = a==b
-
-        Color c1 = [11,22,33]
-        Color c2 = [11,22,33]
-        Color c3 = [11,22,33]
-        uword[] colors = [ c1, c2, c3]      ; TODO should contain pointers to (the first element) of each struct
+        txt.print_uwhex(&c1, true)
+        txt.chrout('\n')
+        txt.print_uwhex(&c2, true)
+        txt.chrout('\n')
+        txt.print_uwhex(&c3, true)
+        txt.chrout('\n')
+        txt.print_uwhex(colors[0], true)
+        txt.chrout('\n')
+        txt.print_uwhex(colors[1], true)
+        txt.chrout('\n')
+        txt.print_uwhex(colors[2], true)
+        txt.chrout('\n')
 
         c1 = c2
-        ; c1 = [11,22,33]         ; TODO implement rewrite into individual struct member assignments
+        c1 = [11,22,33]         ; TODO implement rewrite into individual struct member assignments
+        testX()
     }
 
     asmsub testX() {