implemented const folding of lmh() and divmod()

codeOptimizers/src/prog8/optimizer/ConstantFoldingOptimizer.kt

@@ -4,10 +4,7 @@ import prog8.ast.*
 import prog8.ast.expressions.*
 import prog8.ast.statements.*
 import prog8.ast.walk.*
-import prog8.code.core.AssociativeOperators
-import prog8.code.core.BaseDataType
-import prog8.code.core.IErrorReporter
-import prog8.code.core.isInteger
+import prog8.code.core.*
 import kotlin.math.floor
 
 
@@ -374,29 +371,33 @@ class ConstantFoldingOptimizer(private val program: Program, private val errors:
     }
 
     override fun after(functionCallExpr: FunctionCallExpression, parent: Node): Iterable<AstModification> {
-        if(functionCallExpr.target.nameInSource == listOf("lmh")) {
-            // builtin that returns 3 values instead of 1
-            if(functionCallExpr.args.size==1) {
-                val constArg = functionCallExpr.args[0].constValue(program)?.number?.toLong()
-                if(constArg!=null) {
-                    val assign = parent as? Assignment
-                    if(assign!=null && assign.target.multi?.size==3) {
-                        val low = constArg and 255
-                        val mid = (constArg shr 8) and 255
-                        val high = (constArg shr 16) and 255
-                        val assignLow = Assignment(assign.target.multi!![0], NumericLiteral(BaseDataType.UBYTE, low.toDouble(), assign.position),AssignmentOrigin.USERCODE, assign.position)
-                        val assignMid = Assignment(assign.target.multi!![1], NumericLiteral(BaseDataType.UBYTE, mid.toDouble(), assign.position), AssignmentOrigin.USERCODE, assign.position)
-                        val assignHigh = Assignment(assign.target.multi!![2], NumericLiteral(BaseDataType.UBYTE, high.toDouble(), assign.position), AssignmentOrigin.USERCODE, assign.position)
-                        return listOf(
-                            AstInsert.after(assign, assignHigh, assign.parent as IStatementContainer),
-                            AstInsert.after(assign, assignMid, assign.parent as IStatementContainer),
-                            AstInsert.after(assign, assignLow, assign.parent as IStatementContainer),
-                            AstRemove(assign, assign.parent as IStatementContainer)
-                        )
-                    } else {
-                        errors.info("constant lmh expression can be replaced by 3 separate constant byte values", functionCallExpr.position)
+        val constValues = functionCallExpr.constValues(program)
+        if (constValues != null) {
+            if (constValues.size == 1) {
+                return listOf(AstReplaceNode(functionCallExpr, constValues[0], parent))
+            } else {
+                // Multi-value return
+                val container = parent.parent as? IStatementContainer ?: return noModifications
+                if (parent is Assignment && parent.target.multi?.size == constValues.size) {
+                    val mods = mutableListOf<AstModification>()
+                    for (i in constValues.indices.reversed()) {
+                        val assign = Assignment(parent.target.multi!![i], constValues[i], AssignmentOrigin.USERCODE, parent.position)
+                        mods.add(AstInsert.after(parent, assign, container))
                     }
-                    return noModifications
+                    mods.add(AstRemove(parent, container))
+                    return mods
+                } else if (parent is VarDecl && parent.names.size == constValues.size) {
+                    val mods = mutableListOf<AstModification>()
+                    for (i in constValues.indices.reversed()) {
+                        val newDecl = VarDecl.builder(DataType.forDt(constValues[i].type), parent.position)
+                            .copyFrom(parent)
+                            .names(parent.names[i])
+                            .value(constValues[i])
+                            .build()
+                        mods.add(AstInsert.after(parent, newDecl, container))
+                    }
+                    mods.add(AstRemove(parent, container))
+                    return mods
                 }
             }
         }

compiler/src/prog8/compiler/BuiltinFunctions.kt

@@ -13,59 +13,60 @@ import prog8.code.core.Position
 import prog8.code.core.isInteger
 import kotlin.math.*
 
-// NOTE: at this time this machinery only supports 1 return value
-private typealias ConstExpressionCaller = (args: List<Expression>, position: Position, program: Program) -> NumericLiteral
-
 // Maps builtin functions to constant evaluators for compile-time execution
-// TODO: add support for builtin functions that return multiple return values (such as divmod and lmh)
-internal val constEvaluatorsForBuiltinFuncs: Map<String, ConstExpressionCaller> = mapOf(
-    "abs__byte" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, true) { abs(it).toDouble() } },
-    "abs__word" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, true) { abs(it).toDouble() } },
-    "abs__long" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, true) { abs(it).toDouble() } },
-    "abs__float" to { a, p, prg -> oneFloatArgOutputFloat(a, p, prg) { abs(it) } },
-    "len" to ::builtinLen,
-    "sizeof" to ::builtinSizeof,
-    "offsetof" to ::builtinOffsetof,
-    "sgn" to ::builtinSgn,
-    "sqrt__ubyte" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, false) { sqrt(it.toDouble()) } },
-    "sqrt__uword" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, false) { sqrt(it.toDouble()) } },
-    "sqrt__long" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, false) {
+internal val constEvaluatorsForBuiltinFuncs: Map<String, (args: List<Expression>, position: Position, program: Program) -> List<NumericLiteral>> = mapOf(
+    "abs__byte" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, true) { abs(it).toDouble() }) },
+    "abs__word" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, true) { abs(it).toDouble() }) },
+    "abs__long" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, true) { abs(it).toDouble() }) },
+    "abs__float" to { a, p, prg -> listOf(oneFloatArgOutputFloat(a, p, prg) { abs(it) }) },
+    "len" to { a, p, prg -> listOf(builtinLen(a, p, prg)) },
+    "sizeof" to { a, p, prg -> listOf(builtinSizeof(a, p, prg)) },
+    "offsetof" to { a, p, prg -> listOf(builtinOffsetof(a, p, prg)) },
+    "sgn" to { a, p, prg -> listOf(builtinSgn(a, p, prg)) },
+    "sqrt__ubyte" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, false) { sqrt(it.toDouble()) }) },
+    "sqrt__uword" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, false) { sqrt(it.toDouble()) }) },
+    "sqrt__long" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, false) {
         val value=it.toDouble()
         if(value<0)
             throw CannotEvaluateException("sqrt", "argument cannot be negative")
         else
             sqrt(value)
-    } },
-    "sqrt__float" to { a, p, prg -> oneFloatArgOutputFloat(a, p, prg) {
+    }) },
+    "sqrt__float" to { a, p, prg -> listOf(oneFloatArgOutputFloat(a, p, prg) {
         if(it<0)
             throw CannotEvaluateException("sqrt", "argument cannot be negative")
         else
             sqrt(it)
-    } },
-    "lsb" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x and 255).toDouble() } },
-    "lsb__long" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x and 255).toDouble() } },
-    "lsw" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x and 65535).toDouble() } },
-    "msb" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x ushr 8 and 255).toDouble()} },
-    "msb__long" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x ushr 24 and 255).toDouble()} },
-    "msw" to { a, p, prg -> oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x ushr 16 and 65535).toDouble()} },
-    "mkword" to ::builtinMkword,
-    "mklong" to ::builtinMklong,
-    "mklong2" to ::builtinMklong,
-    "clamp__ubyte" to ::builtinClampUByte,
-    "clamp__byte" to ::builtinClampByte,
-    "clamp__uword" to ::builtinClampUWord,
-    "clamp__word" to ::builtinClampWord,
-    "clamp__long" to ::builtinClampLong,
-    "min__ubyte" to ::builtinMinUByte,
-    "min__byte" to ::builtinMinByte,
-    "min__uword" to ::builtinMinUWord,
-    "min__word" to ::builtinMinWord,
-    "min__long" to ::builtinMinLong,
-    "max__ubyte" to ::builtinMaxUByte,
-    "max__byte" to ::builtinMaxByte,
-    "max__uword" to ::builtinMaxUWord,
-    "max__word" to ::builtinMaxWord,
-    "max__long" to ::builtinMaxLong
+    }) },
+    "lsb" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x and 255).toDouble() }) },
+    "lsb__long" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x and 255).toDouble() }) },
+    "lsw" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x and 65535).toDouble() }) },
+    "msb" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x ushr 8 and 255).toDouble()}) },
+    "msb__long" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x ushr 24 and 255).toDouble()}) },
+    "msw" to { a, p, prg -> listOf(oneIntArgOutputInt(a, p, prg, true) { x: Int -> (x ushr 16 and 65535).toDouble()}) },
+    "mkword" to { a, p, prg -> listOf(builtinMkword(a, p, prg)) },
+    "mklong" to { a, p, prg -> listOf(builtinMklong(a, p, prg)) },
+    "mklong2" to { a, p, prg -> listOf(builtinMklong(a, p, prg)) },
+    "clamp__ubyte" to { a, p, prg -> listOf(builtinClampUByte(a, p, prg)) },
+    "clamp__byte" to { a, p, prg -> listOf(builtinClampByte(a, p, prg)) },
+    "clamp__uword" to { a, p, prg -> listOf(builtinClampUWord(a, p, prg)) },
+    "clamp__word" to { a, p, prg -> listOf(builtinClampWord(a, p, prg)) },
+    "clamp__long" to { a, p, prg -> listOf(builtinClampLong(a, p, prg)) },
+    "min__ubyte" to { a, p, prg -> listOf(builtinMinUByte(a, p, prg)) },
+    "min__byte" to { a, p, prg -> listOf(builtinMinByte(a, p, prg)) },
+    "min__uword" to { a, p, prg -> listOf(builtinMinUWord(a, p, prg)) },
+    "min__word" to { a, p, prg -> listOf(builtinMinWord(a, p, prg)) },
+    "min__long" to { a, p, prg -> listOf(builtinMinLong(a, p, prg)) },
+    "max__ubyte" to { a, p, prg -> listOf(builtinMaxUByte(a, p, prg)) },
+    "max__byte" to { a, p, prg -> listOf(builtinMaxByte(a, p, prg)) },
+    "max__uword" to { a, p, prg -> listOf(builtinMaxUWord(a, p, prg)) },
+    "max__word" to { a, p, prg -> listOf(builtinMaxWord(a, p, prg)) },
+    "max__long" to { a, p, prg -> listOf(builtinMaxLong(a, p, prg)) },
+    "lmh" to ::builtinLmh,
+    "divmod__ubyte" to ::builtinDivmodUByte,
+    "divmod__uword" to ::builtinDivmodUWord,
+    "divmod__byte" to ::builtinDivmodByte,
+    "divmod__word" to ::builtinDivmodWord
 )
 
 internal fun builtinFunctionReturnTypes(function: String): Array<InferredTypes.InferredType> {
@@ -397,3 +398,63 @@ private fun builtinClampLong(args: List<Expression>, position: Position, program
     return NumericLiteral(BaseDataType.LONG, result, position)
 }
 
+
+private fun builtinLmh(args: List<Expression>, position: Position, program: Program): List<NumericLiteral> {
+    if (args.size != 1)
+        throw SyntaxError("lmh requires 1 argument", position)
+    val constval = args[0].constValue(program) ?: throw NotConstArgumentException()
+    val value = constval.number.toLong()
+    return listOf(
+        NumericLiteral(BaseDataType.UBYTE, (value and 0xff).toDouble(), position),
+        NumericLiteral(BaseDataType.UBYTE, ((value shr 8) and 0xff).toDouble(), position),
+        NumericLiteral(BaseDataType.UBYTE, ((value shr 16) and 0xff).toDouble(), position)
+    )
+}
+
+private fun builtinDivmodUByte(args: List<Expression>, position: Position, program: Program): List<NumericLiteral> {
+    if (args.size != 2)
+        throw SyntaxError("divmod requires 2 arguments", position)
+    val v1 = args[0].constValue(program) ?: throw NotConstArgumentException()
+    val v2 = args[1].constValue(program) ?: throw NotConstArgumentException()
+    if (v2.number.toInt() == 0) throw CannotEvaluateException("divmod", "division by zero")
+    return listOf(
+        NumericLiteral(BaseDataType.UBYTE, (v1.number.toInt() / v2.number.toInt()).toDouble(), position),
+        NumericLiteral(BaseDataType.UBYTE, (v1.number.toInt() % v2.number.toInt()).toDouble(), position)
+    )
+}
+
+private fun builtinDivmodUWord(args: List<Expression>, position: Position, program: Program): List<NumericLiteral> {
+    if (args.size != 2)
+        throw SyntaxError("divmod requires 2 arguments", position)
+    val v1 = args[0].constValue(program) ?: throw NotConstArgumentException()
+    val v2 = args[1].constValue(program) ?: throw NotConstArgumentException()
+    if (v2.number.toInt() == 0) throw CannotEvaluateException("divmod", "division by zero")
+    return listOf(
+        NumericLiteral(BaseDataType.UWORD, (v1.number.toInt() / v2.number.toInt()).toDouble(), position),
+        NumericLiteral(BaseDataType.UWORD, (v1.number.toInt() % v2.number.toInt()).toDouble(), position)
+    )
+}
+
+private fun builtinDivmodByte(args: List<Expression>, position: Position, program: Program): List<NumericLiteral> {
+    if (args.size != 2)
+        throw SyntaxError("divmod requires 2 arguments", position)
+    val v1 = args[0].constValue(program) ?: throw NotConstArgumentException()
+    val v2 = args[1].constValue(program) ?: throw NotConstArgumentException()
+    if (v2.number.toInt() == 0) throw CannotEvaluateException("divmod", "division by zero")
+    return listOf(
+        NumericLiteral(BaseDataType.BYTE, (v1.number.toInt() / v2.number.toInt()).toDouble(), position),
+        NumericLiteral(BaseDataType.BYTE, (v1.number.toInt() % v2.number.toInt()).toDouble(), position)
+    )
+}
+
+private fun builtinDivmodWord(args: List<Expression>, position: Position, program: Program): List<NumericLiteral> {
+    if (args.size != 2)
+        throw SyntaxError("divmod requires 2 arguments", position)
+    val v1 = args[0].constValue(program) ?: throw NotConstArgumentException()
+    val v2 = args[1].constValue(program) ?: throw NotConstArgumentException()
+    if (v2.number.toInt() == 0) throw CannotEvaluateException("divmod", "division by zero")
+    return listOf(
+        NumericLiteral(BaseDataType.WORD, (v1.number.toInt() / v2.number.toInt()).toDouble(), position),
+        NumericLiteral(BaseDataType.WORD, (v1.number.toInt() % v2.number.toInt()).toDouble(), position)
+    )
+}

compiler/src/prog8/compiler/Compiler.kt

@@ -376,9 +376,9 @@ private class BuiltinFunctionsFacade(functions: Map<String, FSignature>): IBuilt
     override val names = functions.keys
     override val purefunctionNames = functions.filter { it.value.pure }.mapTo(mutableSetOf()) { it.key }
 
-    override fun constValue(funcName: String, args: List<Expression>, position: Position): NumericLiteral? {
+    override fun constValues(funcName: String, args: List<Expression>, position: Position): List<NumericLiteral>? {
         if(funcName=="msw" && !args[0].inferType(program).isLong)
-            return NumericLiteral.optimalInteger(0, position)
+            return listOf(NumericLiteral.optimalInteger(0, position))
 
         val func = BuiltinFunctions[funcName]
         if(func!=null) {

compiler/test/TestBuiltinFunctions.kt

@@ -6,6 +6,7 @@ import io.kotest.inspectors.shouldForAll
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.shouldNotBe
 import io.kotest.matchers.string.shouldContain
+import io.kotest.matchers.string.shouldNotContain
 import prog8.ast.expressions.NumericLiteral
 import prog8.ast.statements.Assignment
 import prog8.ast.statements.FunctionCallStatement
@@ -171,6 +172,37 @@ main {
         compileText(C64Target(), false, src, outputDir, writeAssembly = true) shouldNotBe null
     }
 
+    test("divmod and lmh should be compile time folded") {
+        val src="""
+main {
+    ubyte @shared d, r, l, m, h
+    sub start() {
+        d, r = divmod(100, 13)
+        l, m, h = lmh(123456)
+    }
+}"""
+        val result = compileText(C64Target(), true, src, outputDir, writeAssembly = false)!!
+        val statements = result.compilerAst.entrypoint.statements
+        
+        val assignments = statements.filterIsInstance<Assignment>()
+        val d = assignments.find { it.target.identifier?.nameInSource?.last() == "d" }!!
+        val r = assignments.find { it.target.identifier?.nameInSource?.last() == "r" }!!
+        val l = assignments.find { it.target.identifier?.nameInSource?.last() == "l" }!!
+        val m = assignments.find { it.target.identifier?.nameInSource?.last() == "m" }!!
+        val h = assignments.find { it.target.identifier?.nameInSource?.last() == "h" }!!
+
+        (d.value as NumericLiteral).number shouldBe 7.0
+        (r.value as NumericLiteral).number shouldBe 9.0
+        (l.value as NumericLiteral).number shouldBe 64.0
+        (m.value as NumericLiteral).number shouldBe 226.0
+        (h.value as NumericLiteral).number shouldBe 1.0
+
+        statements.forEach { 
+             it.toString() shouldNotContain "divmod("
+             it.toString() shouldNotContain "lmh("
+        }
+    }
+
     test("sgn") {
         val src="""
 main {

compiler/test/helpers/Dummies.kt

@@ -16,11 +16,11 @@ internal object DummyFunctions : IBuiltinFunctions {
     override val names: Set<String> = emptySet()
     override val purefunctionNames: Set<String> = emptySet()
     
-    override fun constValue(
+    override fun constValues(
         funcName: String,
         args: List<Expression>,
         position: Position,
-    ): NumericLiteral? = null
+    ): List<NumericLiteral>? = null
 
     override fun returnTypes(funcName: String) = emptyArray<InferredTypes.InferredType>()
 }

compilerAst/src/prog8/ast/IBuiltinFunctions.kt

@@ -8,6 +8,7 @@ import prog8.code.core.Position
 interface IBuiltinFunctions {
     val names: Set<String>
     val purefunctionNames: Set<String>
-    fun constValue(funcName: String, args: List<Expression>, position: Position): NumericLiteral?
+    fun constValue(funcName: String, args: List<Expression>, position: Position): NumericLiteral? = constValues(funcName, args, position)?.singleOrNull()
+    fun constValues(funcName: String, args: List<Expression>, position: Position): List<NumericLiteral>?
     fun returnTypes(funcName: String): Array<InferredTypes.InferredType>
 }

compilerAst/src/prog8/ast/expressions/AstExpressions.kt

@@ -17,6 +17,10 @@ import kotlin.math.truncate
 sealed class Expression: Node {
     abstract override fun copy(): Expression
     abstract fun constValue(program: Program): NumericLiteral?
+    open fun constValues(program: Program): List<NumericLiteral>? {
+        val cv = constValue(program)
+        return if(cv!=null) listOf(cv) else null
+    }
     abstract fun accept(visitor: IAstVisitor)
     abstract fun accept(visitor: AstWalker, parent: Node)
     abstract fun inferType(program: Program): InferredTypes.InferredType
@@ -1634,6 +1638,14 @@ class FunctionCallExpression(override var target: IdentifierReference,
 
     override fun constValue(program: Program) = constValue(program, true)
 
+    override fun constValues(program: Program): List<NumericLiteral>? {
+        if (target.nameInSource.size > 1)
+            return null
+        val values = program.builtinFunctions.constValues(target.nameInSource[0], args, position)
+        values?.forEach { it.parent = this.parent }
+        return values
+    }
+
     private fun constValue(program: Program, withDatatypeCheck: Boolean): NumericLiteral? {
         // if the function is a built-in function and the args are consts, should try to const-evaluate!
         // lenghts of arrays and strings are constants that are determined at compile time!

examples/test.p8

@@ -4,9 +4,9 @@
 main {
 
     sub start() {
-        long @shared lv = $12345678
+        const long lv = $12345678
 
-        ubyte l,m,h = lmh(lv+99999)
+        ubyte l,m,h = lmh(lv+99999)     ; TODO should be const-folded and ALL THREE target vars should become consts too
 
         txt.print_ubhex(h, true)
         txt.spc()
@@ -16,10 +16,10 @@ main {
         txt.nl()
 
 
-        ubyte @shared num = 230
-        ubyte @shared div = 13
+        const ubyte num = 230
+        const ubyte div = 13
 
-        ubyte d,r = divmod(num, div)
+        ubyte d,r = divmod(num, div)     ; TODO should be const-folded and BOTH target vars should become consts too
         txt.print_ub(d)
         txt.spc()
         txt.print_ub(r)