Implement ESTree based code generation for the MST

kmath-ast/build.gradle.kts

@@ -2,6 +2,20 @@ plugins {
     id("ru.mipt.npm.mpp")
 }
 
+kotlin.js {
+    nodejs {
+        testTask {
+            useMocha().timeout = "0"
+        }
+    }
+
+    browser {
+        testTask {
+            useMocha().timeout = "0"
+        }
+    }
+}
+
 kotlin.sourceSets {
     commonMain {
         dependencies {

kmath-ast/src/jsMain/kotlin/astring.global.kt

@@ -10,7 +10,7 @@ import Generator
 
 @Suppress("EXTERNAL_DELEGATION", "NESTED_CLASS_IN_EXTERNAL_INTERFACE")
 external interface astring {
-    var generate: Any
+    var generate: (dynamic, dynamic) -> dynamic
     var baseGenerator: Generator
 
     companion object : astring by definedExternally

kmath-ast/src/jsMain/kotlin/astring.kt

@@ -1,7 +1,15 @@
-@file:Suppress("INTERFACE_WITH_SUPERCLASS", "OVERRIDING_FINAL_MEMBER", "RETURN_TYPE_MISMATCH_ON_OVERRIDE", "CONFLICTING_OVERLOADS",
+@file:Suppress(
-    "NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "KDocMissingDocumentation"
+    "INTERFACE_WITH_SUPERCLASS", "OVERRIDING_FINAL_MEMBER", "RETURN_TYPE_MISMATCH_ON_OVERRIDE", "CONFLICTING_OVERLOADS",
+    "NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "KDocMissingDocumentation", "PackageDirectoryMismatch"
 )
 
+@file:JsModule("astring")
+@file:JsNonModule
+package astring
+
+import Generator
+import estree.BaseNode
+
 external interface Options {
     var indent: String?
         get() = definedExternally
@@ -27,6 +35,4 @@ external fun generate(node: BaseNode, options: Options /* Options & `T$0` */ = d
 
 external fun generate(node: BaseNode): String
 
-typealias Generator = Any
+external var baseGenerator: Generator
-
-external var baseGenerator: Generator

kmath-ast/src/jsMain/kotlin/astring.typealises.kt

@@ -0,0 +1,3 @@
+@file:Suppress("NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "KDocMissingDocumentation")
+
+typealias Generator = Any

kmath-ast/src/jsMain/kotlin/emitter.d.ts

@@ -14,6 +14,3 @@ export class Emitter {
 
     hasListeners(event: string): boolean
 }
-
-
-function mixin(obj: any): any

kmath-ast/src/jsMain/kotlin/emitter.kt

@@ -2,9 +2,12 @@
     "INTERFACE_WITH_SUPERCLASS",
     "OVERRIDING_FINAL_MEMBER",
     "RETURN_TYPE_MISMATCH_ON_OVERRIDE",
-    "CONFLICTING_OVERLOADS", "NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "KDocMissingDocumentation", "SortModifiers"
+    "CONFLICTING_OVERLOADS", "NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "KDocMissingDocumentation", "SortModifiers",
+    "PackageDirectoryMismatch"
 )
 
+package emitter
+
 external open class Emitter {
     constructor(obj: Any)
     constructor()

kmath-ast/src/jsMain/kotlin/estree.extensions.kt

@@ -0,0 +1,67 @@
+@file:Suppress(
+    "NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "KDocMissingDocumentation",
+    "NO_EXPLICIT_RETURN_TYPE_IN_API_MODE_WARNING", "PackageDirectoryMismatch"
+)
+
+package estree
+
+fun Program(sourceType: String, vararg body: dynamic) = object : Program {
+    override var type = "Program"
+    override var sourceType = sourceType
+    override var body = body
+}
+
+fun VariableDeclaration(kind: String, vararg declarations: VariableDeclarator) = object : VariableDeclaration {
+    override var type = "VariableDeclaration"
+    override var declarations = declarations.toList().toTypedArray()
+    override var kind = kind
+}
+
+fun VariableDeclarator(id: dynamic, init: dynamic) = object : VariableDeclarator {
+    override var type = "VariableDeclarator"
+    override var id = id
+    override var init = init
+}
+
+fun Identifier(name: String) = object : Identifier {
+    override var type = "Identifier"
+    override var name = name
+}
+
+fun FunctionExpression(params: Array<dynamic>, body: BlockStatement) = object : FunctionExpression {
+    override var params = params
+    override var type = "FunctionExpression"
+    override var body = body
+}
+
+fun BlockStatement(vararg body: dynamic) = object : BlockStatement {
+    override var type = "BlockStatement"
+    override var body = body
+}
+
+fun ReturnStatement(argument: dynamic) = object : ReturnStatement {
+    override var type = "ReturnStatement"
+    override var argument = argument
+}
+
+fun SimpleLiteral(value: dynamic) = object : SimpleLiteral {
+    override var type = "Literal"
+    override var value = value
+}
+
+fun MemberExpression(computed: Boolean, optional: Boolean, `object`: dynamic, property: dynamic) =
+    object : MemberExpression {
+        override var type = "MemberExpression"
+        override var computed = computed
+        override var optional = optional
+        override var `object` = `object`
+        override var property = property
+    }
+
+fun SimpleCallExpression(optional: Boolean, callee: dynamic, vararg arguments: dynamic) =
+    object : SimpleCallExpression {
+        override var type = "CallExpression"
+        override var optional = optional
+        override var callee = callee
+        override var arguments = arguments
+    }

kmath-ast/src/jsMain/kotlin/estree.kt

@@ -1,8 +1,10 @@
 @file:Suppress(
     "INTERFACE_WITH_SUPERCLASS", "OVERRIDING_FINAL_MEMBER", "RETURN_TYPE_MISMATCH_ON_OVERRIDE", "CONFLICTING_OVERLOADS",
-    "NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "KDocMissingDocumentation", "ClassName",
+    "NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "KDocMissingDocumentation", "ClassName", "PackageDirectoryMismatch",
 )
 
+package estree
+
 import kotlin.js.RegExp
 
 external interface BaseNodeWithoutComments {

kmath-ast/src/jsMain/kotlin/kscience/kmath/estree/Codegen.kt

@@ -0,0 +1,78 @@
+package kscience.kmath.estree
+
+import estree.*
+import kscience.kmath.ast.MST
+import kscience.kmath.ast.MstExpression
+import kscience.kmath.estree.internal.JSBuilder
+import kscience.kmath.expressions.Expression
+import kscience.kmath.operations.Algebra
+import kscience.kmath.operations.NumericAlgebra
+import kscience.kmath.operations.RealField
+
+@PublishedApi
+internal fun <T> MST.compileWith(algebra: Algebra<T>): Expression<T> {
+    fun JSBuilder<T>.visit(node: MST): BaseExpression = when (node) {
+        is MST.Symbolic -> {
+            val symbol = try {
+                algebra.symbol(node.value)
+            } catch (ignored: IllegalStateException) {
+                null
+            }
+
+            if (symbol != null)
+                constant(symbol)
+            else
+                variable(node.value)
+        }
+
+        is MST.Numeric -> constant(node.value)
+        is MST.Unary -> call(algebra.unaryOperation(node.operation), visit(node.value))
+
+        is MST.Binary -> when {
+            algebra is NumericAlgebra<T> && node.left is MST.Numeric && node.right is MST.Numeric -> constant(
+                algebra.number(
+                    RealField
+                        .binaryOperation(node.operation)
+                        .invoke(node.left.value.toDouble(), node.right.value.toDouble())
+                )
+            )
+
+            algebra is NumericAlgebra<T> && node.left is MST.Numeric -> call(
+                algebra.leftSideNumberOperation(node.operation),
+                visit(node.left),
+                visit(node.right),
+            )
+
+            algebra is NumericAlgebra<T> && node.right is MST.Numeric -> call(
+                algebra.rightSideNumberOperation(node.operation),
+                visit(node.left),
+                visit(node.right),
+            )
+
+            else -> call(
+                algebra.binaryOperation(node.operation),
+                visit(node.left),
+                visit(node.right),
+            )
+        }
+    }
+
+    return JSBuilder<T> { visit(this@compileWith) }.instance
+}
+
+
+/**
+ * Compiles an [MST] to ASM using given algebra.
+ *
+ * @author Alexander Nozik.
+ */
+public fun <T : Any> Algebra<T>.expression(mst: MST): Expression<T> =
+    mst.compileWith(this)
+
+/**
+ * Optimizes performance of an [MstExpression] using ASM codegen.
+ *
+ * @author Alexander Nozik.
+ */
+public fun <T : Any> MstExpression<T, Algebra<T>>.compile(): Expression<T> =
+    mst.compileWith(algebra)

kmath-ast/src/jsMain/kotlin/kscience/kmath/estree/internal/JSBuilder.kt

@@ -0,0 +1,73 @@
+package kscience.kmath.estree.internal
+
+import astring.generate
+import estree.*
+import kscience.kmath.expressions.Expression
+import kscience.kmath.expressions.Symbol
+
+internal class JSBuilder<T>(val bodyCallback: JSBuilder<T>.() -> BaseExpression) {
+    private class GeneratedExpression<T>(val executable: dynamic, val constants: Array<dynamic>) : Expression<T> {
+        @Suppress("UNUSED_VARIABLE", "PARAMETER_NAME_CHANGED_ON_OVERRIDE")
+        override fun invoke(map: Map<Symbol, T>): T {
+            val e = executable
+            val c = constants
+            val a = js("{}")
+            map.forEach { (key, value) -> a[key.identity] = value }
+            return js("e(c, a)").unsafeCast<T>()
+        }
+    }
+
+    val instance: Expression<T> by lazy {
+        val node = Program(
+            sourceType = "script",
+            VariableDeclaration(
+                kind = "var",
+                VariableDeclarator(
+                    id = Identifier("executable"),
+                    init = FunctionExpression(
+                        params = arrayOf(Identifier("constants"), Identifier("arguments")),
+                        body = BlockStatement(ReturnStatement(bodyCallback())),
+                    ),
+                ),
+            ),
+        )
+
+        eval(generate(node))
+        GeneratedExpression(js("executable"), constants.toTypedArray())
+    }
+
+    private val constants = mutableListOf<Any>()
+    private val keys = mutableListOf<String>()
+
+    fun constant(value: Any?) = when {
+        value == null || jsTypeOf(value) == "number" || jsTypeOf(value) == "string" || jsTypeOf(value) == "boolean" -> {
+            SimpleLiteral(value)
+        }
+
+        else -> {
+            val idx = if (value in constants) constants.indexOf(value) else constants.also { it += value }.lastIndex
+
+            MemberExpression(
+                computed = true,
+                optional = false,
+                `object` = Identifier("constants"),я
+                property = SimpleLiteral(idx),
+            )
+        }
+    }
+
+    fun variable(name: String): BaseExpression {
+        return MemberExpression(
+            computed = true,
+            optional = false,
+            `object` = Identifier("arguments"),
+            property = SimpleLiteral(name),
+        )
+    }
+
+    fun call(function: Function<T>, vararg args: BaseExpression): BaseExpression = SimpleCallExpression(
+        optional = false,
+        callee = constant(function),
+        *args,
+    )
+}

kmath-ast/src/jsMain/kotlin/lib.es2015.iterable.kt

@@ -1,20 +1,8 @@
-@file:Suppress("INTERFACE_WITH_SUPERCLASS", "OVERRIDING_FINAL_MEMBER", "RETURN_TYPE_MISMATCH_ON_OVERRIDE", "CONFLICTING_OVERLOADS")
+@file:Suppress("INTERFACE_WITH_SUPERCLASS", "OVERRIDING_FINAL_MEMBER", "RETURN_TYPE_MISMATCH_ON_OVERRIDE", "CONFLICTING_OVERLOADS",
+    "NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "KDocMissingDocumentation", "PackageDirectoryMismatch"
+)
 package tsstdlib
 
-import kotlin.js.*
-import org.khronos.webgl.*
-import org.w3c.dom.*
-import org.w3c.dom.events.*
-import org.w3c.dom.parsing.*
-import org.w3c.dom.svg.*
-import org.w3c.dom.url.*
-import org.w3c.fetch.*
-import org.w3c.files.*
-import org.w3c.notifications.*
-import org.w3c.performance.*
-import org.w3c.workers.*
-import org.w3c.xhr.*
-
 external interface IteratorYieldResult<TYield> {
     var done: Boolean?
         get() = definedExternally

kmath-ast/src/jsMain/kotlin/stream.kt

@@ -1,9 +1,13 @@
 @file:Suppress(
     "INTERFACE_WITH_SUPERCLASS", "OVERRIDING_FINAL_MEMBER", "RETURN_TYPE_MISMATCH_ON_OVERRIDE", "CONFLICTING_OVERLOADS",
     "NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "NO_EXPLICIT_VISIBILITY_IN_API_MODE_WARNING", "SortModifiers",
-    "KDocMissingDocumentation"
+    "KDocMissingDocumentation", "PackageDirectoryMismatch"
 )
 
+package stream
+
+import emitter.Emitter
+
 external open class Stream : Emitter {
     open fun pipe(dest: Any, options: Any): Any
 }

kmath-ast/src/jsTest/kotlin/kscience/kmath/estree/TestESTreeAlgebras.kt

@@ -0,0 +1,94 @@
+package kscience.kmath.estree
+
+import kscience.kmath.ast.mstInField
+import kscience.kmath.ast.mstInRing
+import kscience.kmath.ast.mstInSpace
+import kscience.kmath.expressions.invoke
+import kscience.kmath.operations.ByteRing
+import kscience.kmath.operations.RealField
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+internal class TestESTreeAlgebras {
+    @Test
+    fun space() {
+        val res1 = ByteRing.mstInSpace {
+            binaryOperation("+")(
+                unaryOperation("+")(
+                    number(3.toByte()) - (number(2.toByte()) + (multiply(
+                        add(number(1), number(1)),
+                        2
+                    ) + number(1.toByte()) * 3.toByte() - number(1.toByte())))
+                ),
+
+                number(1)
+            ) + symbol("x") + zero
+        }("x" to 2.toByte())
+
+        val res2 = ByteRing.mstInSpace {
+            binaryOperation("+")(
+                unaryOperation("+")(
+                    number(3.toByte()) - (number(2.toByte()) + (multiply(
+                        add(number(1), number(1)),
+                        2
+                    ) + number(1.toByte()) * 3.toByte() - number(1.toByte())))
+                ),
+
+                number(1)
+            ) + symbol("x") + zero
+        }.compile()("x" to 2.toByte())
+
+        assertEquals(res1, res2)
+    }
+
+    @Test
+    fun ring() {
+        val res1 = ByteRing.mstInRing {
+            binaryOperation("+")(
+                unaryOperation("+")(
+                    (symbol("x") - (2.toByte() + (multiply(
+                        add(number(1), number(1)),
+                        2
+                    ) + 1.toByte()))) * 3.0 - 1.toByte()
+                ),
+
+                number(1)
+            ) * number(2)
+        }("x" to 3.toByte())
+
+        val res2 = ByteRing.mstInRing {
+            binaryOperation("+")(
+                unaryOperation("+")(
+                    (symbol("x") - (2.toByte() + (multiply(
+                        add(number(1), number(1)),
+                        2
+                    ) + 1.toByte()))) * 3.0 - 1.toByte()
+                ),
+                number(1)
+            ) * number(2)
+        }.compile()("x" to 3.toByte())
+
+        assertEquals(res1, res2)
+    }
+
+    @Test
+    fun field() {
+        val res1 = RealField.mstInField {
+            +(3 - 2 + 2 * number(1) + 1.0) + binaryOperation("+")(
+                (3.0 - (symbol("x") + (multiply(add(number(1.0), number(1.0)), 2) + 1.0))) * 3 - 1.0
+                        + number(1),
+                number(1) / 2 + number(2.0) * one
+            ) + zero
+        }("x" to 2.0)
+
+        val res2 = RealField.mstInField {
+            +(3 - 2 + 2 * number(1) + 1.0) + binaryOperation("+")(
+                (3.0 - (symbol("x") + (multiply(add(number(1.0), number(1.0)), 2) + 1.0))) * 3 - 1.0
+                        + number(1),
+                number(1) / 2 + number(2.0) * one
+            ) + zero
+        }.compile()("x" to 2.0)
+
+        assertEquals(res1, res2)
+    }
+}

kmath-ast/src/jsTest/kotlin/kscience/kmath/estree/TestESTreeExpressions.kt

@@ -0,0 +1,56 @@
+package kscience.kmath.estree
+
+import kscience.kmath.ast.mstInExtendedField
+import kscience.kmath.ast.mstInField
+import kscience.kmath.ast.mstInSpace
+import kscience.kmath.expressions.Expression
+import kscience.kmath.expressions.StringSymbol
+import kscience.kmath.expressions.invoke
+import kscience.kmath.operations.RealField
+import kotlin.math.pow
+import kotlin.random.Random
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.time.measureTime
+
+internal class TestESTreeExpressions {
+    @Test
+    fun testUnaryOperationInvocation() {
+        val expression = RealField.mstInSpace { -symbol("x") }.compile()
+        val res = expression("x" to 2.0)
+        assertEquals(-2.0, res)
+    }
+
+    @Test
+    fun testBinaryOperationInvocation() {
+        val expression = RealField.mstInSpace { -symbol("x") + number(1.0) }.compile()
+        val res = expression("x" to 2.0)
+        assertEquals(-1.0, res)
+    }
+
+    @Test
+    fun testConstProductInvocation() {
+        val res = RealField.mstInField { symbol("x") * 2 }("x" to 2.0)
+        assertEquals(4.0, res)
+    }
+
+    @Test
+    fun testMultipleCalls() {
+        val e1 =
+            RealField.mstInExtendedField { sin(symbol("x")).pow(4) - 6 * symbol("x") / tanh(symbol("x")) }.compile()
+
+        val e2 = Expression<Double> { a ->
+            val x = a.getValue(StringSymbol("x"))
+            kotlin.math.sin(x).pow(4) - 6 * x / kotlin.math.tanh(x)
+        }
+
+        var r = Random(0)
+        var s = 0.0
+        measureTime { repeat(1000000) { s += e1("x" to r.nextDouble()) } }.also(::println)
+        println(s)
+        s = 0.0
+        r = Random(0)
+        measureTime { repeat(1000000) { s += e2("x" to r.nextDouble()) } }.also(::println)
+        println(s)
+    }
+}

kmath-ast/src/jsTest/kotlin/kscience/kmath/estree/TestESTreeSpecialization.kt

@@ -0,0 +1,54 @@
+package kscience.kmath.estree
+
+import kscience.kmath.ast.mstInField
+import kscience.kmath.expressions.invoke
+import kscience.kmath.operations.RealField
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+internal class TestESTreeSpecialization {
+    @Test
+    fun testUnaryPlus() {
+        val expr = RealField.mstInField { unaryOperation("+")(symbol("x")) }.compile()
+        assertEquals(2.0, expr("x" to 2.0))
+    }
+
+    @Test
+    fun testUnaryMinus() {
+        val expr = RealField.mstInField { unaryOperation("-")(symbol("x")) }.compile()
+        assertEquals(-2.0, expr("x" to 2.0))
+    }
+
+    @Test
+    fun testAdd() {
+        val expr = RealField.mstInField { binaryOperation("+")(symbol("x"), symbol("x")) }.compile()
+        assertEquals(4.0, expr("x" to 2.0))
+    }
+
+    @Test
+    fun testSine() {
+        val expr = RealField.mstInField { unaryOperation("sin")(symbol("x")) }.compile()
+        assertEquals(0.0, expr("x" to 0.0))
+    }
+
+    @Test
+    fun testMinus() {
+        val expr = RealField.mstInField { binaryOperation("-")(symbol("x"), symbol("x")) }.compile()
+        assertEquals(0.0, expr("x" to 2.0))
+    }
+
+    @Test
+    fun testDivide() {
+        val expr = RealField.mstInField { binaryOperation("/")(symbol("x"), symbol("x")) }.compile()
+        assertEquals(1.0, expr("x" to 2.0))
+    }
+
+    @Test
+    fun testPower() {
+        val expr = RealField
+            .mstInField { binaryOperation("pow")(symbol("x"), number(2)) }
+            .compile()
+
+        assertEquals(4.0, expr("x" to 2.0))
+    }
+}

kmath-ast/src/jsTest/kotlin/kscience/kmath/estree/TestESTreeVariables.kt

@@ -0,0 +1,22 @@
+package kscience.kmath.estree
+
+import kscience.kmath.ast.mstInRing
+import kscience.kmath.expressions.invoke
+import kscience.kmath.operations.ByteRing
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+
+internal class TestESTreeVariables {
+    @Test
+    fun testVariableWithoutDefault() {
+        val expr = ByteRing.mstInRing { symbol("x") }.compile()
+        assertEquals(1.toByte(), expr("x" to 1.toByte()))
+    }
+
+    @Test
+    fun testVariableWithoutDefaultFails() {
+        val expr = ByteRing.mstInRing { symbol("x") }.compile()
+        assertFailsWith<NoSuchElementException> { expr() }
+    }
+}

kmath-ast/src/jvmMain/kotlin/kscience/kmath/asm/internal/AsmBuilder.kt

@@ -274,7 +274,7 @@ internal class AsmBuilder<T>(
 
     inline fun buildCall(function: Function<T>, parameters: AsmBuilder<T>.() -> Unit) {
         contract { callsInPlace(parameters, InvocationKind.EXACTLY_ONCE) }
-        val `interface` = function.javaClass.interfaces.first { it.interfaces.contains(Function::class.java) }
+        val `interface` = function.javaClass.interfaces.first { Function::class.java in it.interfaces }
 
         val arity = `interface`.methods.find { it.name == "invoke" }?.parameterCount
             ?: error("Provided function object doesn't contain invoke method")

kmath-ast/src/jvmTest/kotlin/kscience/kmath/asm/TestAsmVariables.kt

@@ -10,13 +10,13 @@ import kotlin.test.assertFailsWith
 internal class TestAsmVariables {
     @Test
     fun testVariableWithoutDefault() {
-        val expr = ByteRing.mstInRing { symbol("x") }
+        val expr = ByteRing.mstInRing { symbol("x") }.compile()
         assertEquals(1.toByte(), expr("x" to 1.toByte()))
     }
 
     @Test
     fun testVariableWithoutDefaultFails() {
-        val expr = ByteRing.mstInRing { symbol("x") }
+        val expr = ByteRing.mstInRing { symbol("x") }.compile()
         assertFailsWith<NoSuchElementException> { expr() }
     }
 }