From 6c1c49d15e17106f385df3ef5eb04e4e581371fc Mon Sep 17 00:00:00 2001
From: Alexander Nozik <altavir@gmail.com>
Date: Sat, 14 Sep 2019 16:50:25 +0300
Subject: [PATCH] Added strong typing to tasks and task dependencies

---
 .../hep/dataforge/workspace/Dependency.kt     |   6 +-
 .../hep/dataforge/workspace/GenericTask.kt    |   2 +-
 .../dataforge/workspace/SimpleWorkspace.kt    |   4 +-
 .../hep/dataforge/workspace/TaskBuilder.kt    |  52 +++++---
 .../hep/dataforge/workspace/TaskModel.kt      | 112 +++++++++++-------
 .../hep/dataforge/workspace/Workspace.kt      |   7 ++
 .../dataforge/workspace/WorkspaceBuilder.kt   |   5 +-
 .../hep/dataforge/workspace/taskTemplates.kt  |   5 -
 .../workspace/SimpleWorkspaceTest.kt          |  36 ++----
 9 files changed, 135 insertions(+), 94 deletions(-)
 delete mode 100644 dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/taskTemplates.kt

diff --git a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/Dependency.kt b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/Dependency.kt
index 1f184095..33b7c8d2 100644
--- a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/Dependency.kt
+++ b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/Dependency.kt
@@ -87,7 +87,11 @@ class DirectTaskDependency<T : Any>(
     }
 }
 
-class WorkspaceTaskDependency(override val name: Name, meta: Meta, placement: Name) : TaskDependency<Any>(meta, placement) {
+class WorkspaceTaskDependency(
+    override val name: Name,
+    meta: Meta,
+    placement: Name
+) : TaskDependency<Any>(meta, placement) {
     override fun resolveTask(workspace: Workspace): Task<*> =
         workspace.tasks[name] ?: error("Task with name $name is not found in the workspace")
 }
\ No newline at end of file
diff --git a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/GenericTask.kt b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/GenericTask.kt
index 408f1745..f59c3d24 100644
--- a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/GenericTask.kt
+++ b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/GenericTask.kt
@@ -55,7 +55,7 @@ class GenericTask<R : Any>(
     override fun build(workspace: Workspace, taskConfig: Meta): TaskModel {
         val taskMeta = taskConfig[name]?.node ?: taskConfig
         val builder = TaskModelBuilder(name, taskMeta)
-        modelTransform.invoke(builder, taskMeta)
+        builder.modelTransform(taskMeta)
         return builder.build()
     }
     //TODO add validation
diff --git a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/SimpleWorkspace.kt b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/SimpleWorkspace.kt
index cb621963..c94fd8ca 100644
--- a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/SimpleWorkspace.kt
+++ b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/SimpleWorkspace.kt
@@ -1,7 +1,6 @@
 package hep.dataforge.workspace
 
 import hep.dataforge.context.Context
-import hep.dataforge.context.Global
 import hep.dataforge.context.content
 import hep.dataforge.context.toMap
 import hep.dataforge.data.DataNode
@@ -24,7 +23,6 @@ class SimpleWorkspace(
     }
 
     companion object {
-        fun build(parent: Context = Global, block: SimpleWorkspaceBuilder.() -> Unit): SimpleWorkspace =
-            SimpleWorkspaceBuilder(parent).apply(block).build()
+
     }
 }
\ No newline at end of file
diff --git a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/TaskBuilder.kt b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/TaskBuilder.kt
index bdacb711..9c0f6958 100644
--- a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/TaskBuilder.kt
+++ b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/TaskBuilder.kt
@@ -8,13 +8,15 @@ import hep.dataforge.meta.get
 import hep.dataforge.meta.string
 import hep.dataforge.names.EmptyName
 import hep.dataforge.names.Name
-import hep.dataforge.names.asName
+import hep.dataforge.names.isEmpty
 import hep.dataforge.names.toName
+import kotlin.jvm.JvmName
 import kotlin.reflect.KClass
 
 @TaskBuildScope
-class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
+class TaskBuilder<R : Any>(val name: Name, val type: KClass<out R>) {
     private var modelTransform: TaskModelBuilder.(Meta) -> Unit = { allData() }
+//    private val additionalDependencies = HashSet<Dependency>()
     var descriptor: NodeDescriptor? = null
     private val dataTransforms: MutableList<DataTransformation> = ArrayList()
 
@@ -30,12 +32,16 @@ class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
             val localData = if (from.isEmpty()) {
                 node
             } else {
-                node[from.toName()].node ?: return null
+                node[from].node ?: return null
             }
             return transform(workspace.context, model, localData)
         }
     }
 
+//    override fun add(dependency: Dependency) {
+//        additionalDependencies.add(dependency)
+//    }
+
     fun model(modelTransform: TaskModelBuilder.(Meta) -> Unit) {
         this.modelTransform = modelTransform
     }
@@ -43,13 +49,14 @@ class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
     /**
      * Add a transformation on untyped data
      */
-    fun rawTransform(
+    @JvmName("rawTransform")
+    fun transform(
         from: String = "",
         to: String = "",
         block: TaskEnv.(DataNode<*>) -> DataNode<R>
     ) {
         dataTransforms += DataTransformation(from, to) { context, model, data ->
-            val env = TaskEnv(EmptyName, model.meta, context)
+            val env = TaskEnv(EmptyName, model.meta, context, data)
             env.block(data)
         }
     }
@@ -62,7 +69,7 @@ class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
     ) {
         dataTransforms += DataTransformation(from, to) { context, model, data ->
             data.ensureType(inputType)
-            val env = TaskEnv(EmptyName, model.meta, context)
+            val env = TaskEnv(EmptyName, model.meta, context, data)
             env.block(data.cast(inputType))
         }
     }
@@ -88,7 +95,14 @@ class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
         }
     }
 
-    class TaskEnv(val name: Name, val meta: Meta, val context: Context)
+    class TaskEnv(val name: Name, val meta: Meta, val context: Context, val data: DataNode<Any>) {
+        operator fun <T : Any> DirectTaskDependency<T>.invoke(): DataNode<T> = if(placement.isEmpty()){
+            data.cast(task.type)
+        } else {
+            data[placement].node?.cast(task.type)
+                ?: error("Could not find results of direct task dependency $this at \"$placement\"")
+        }
+    }
 
     /**
      * A customized pipe action with ability to change meta and name
@@ -121,7 +135,7 @@ class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
             ) {
                 //TODO automatically append task meta
                 result = { data ->
-                    TaskEnv(name, meta, context).block(data)
+                    block(data)
                 }
             }
         }
@@ -133,7 +147,7 @@ class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
     inline fun <reified T : Any> joinByGroup(
         from: String = "",
         to: String = "",
-        crossinline block: JoinGroupBuilder<T, R>.(TaskEnv) -> Unit
+        crossinline block: JoinGroupBuilder<T, R>.(TaskEnv) -> Unit        //TODO needs KEEP-176
     ) {
         action(from, to) {
             JoinAction(
@@ -159,7 +173,7 @@ class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
                     result(
                         actionMeta[TaskModel.MODEL_TARGET_KEY]?.string ?: "@anonymous"
                     ) { data ->
-                        TaskEnv(name, meta, context).block(data)
+                        block(data)
                     }
                 }
             )
@@ -172,7 +186,7 @@ class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
     inline fun <reified T : Any> split(
         from: String = "",
         to: String = "",
-        crossinline block: SplitBuilder<T, R>.(TaskEnv) -> Unit
+        crossinline block: SplitBuilder<T, R>.(TaskEnv) -> Unit  //TODO needs KEEP-176
     ) {
         action(from, to) {
             SplitAction(
@@ -189,9 +203,14 @@ class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
         this.descriptor = NodeDescriptor.build(transform)
     }
 
-    internal fun build(): GenericTask<R> =
-        GenericTask(
-            name.asName(),
+    internal fun build(): GenericTask<R> {
+//        val actualTransform: TaskModelBuilder.(Meta) -> Unit = {
+//            modelTransform
+//            dependencies.addAll(additionalDependencies)
+//        }
+
+        return GenericTask(
+            name,
             type,
             descriptor ?: NodeDescriptor.empty(),
             modelTransform
@@ -220,15 +239,14 @@ class TaskBuilder<R : Any>(val name: String, val type: KClass<out R>) {
                 }
             }
         }
+    }
 }
 
 fun <T : Any> Workspace.Companion.task(
     name: String,
     type: KClass<out T>,
     builder: TaskBuilder<T>.() -> Unit
-): GenericTask<T> {
-    return TaskBuilder(name, type).apply(builder).build()
-}
+): GenericTask<T> = TaskBuilder(name.toName(), type).apply(builder).build()
 
 
 //TODO add delegates to build gradle-like tasks
\ No newline at end of file
diff --git a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/TaskModel.kt b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/TaskModel.kt
index f9562438..7f9cd976 100644
--- a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/TaskModel.kt
+++ b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/TaskModel.kt
@@ -63,58 +63,88 @@ fun TaskModel.buildInput(workspace: Workspace): DataTree<Any> {
 @DslMarker
 annotation class TaskBuildScope
 
+interface TaskDependencyContainer {
+    val defaultMeta: Meta
+    fun add(dependency: Dependency)
+}
+
+/**
+ * Add dependency for a task defined in a workspace and resolved by
+ */
+fun TaskDependencyContainer.dependsOn(
+    name: Name,
+    placement: Name = EmptyName,
+    meta: Meta = defaultMeta
+): WorkspaceTaskDependency =
+    WorkspaceTaskDependency(name, meta, placement).also { add(it) }
+
+fun TaskDependencyContainer.dependsOn(
+    name: String,
+    placement: Name = EmptyName,
+    meta: Meta = defaultMeta
+): WorkspaceTaskDependency =
+    dependsOn(name.toName(), placement, meta)
+
+fun <T : Any> TaskDependencyContainer.dependsOn(
+    task: Task<T>,
+    placement: Name = EmptyName,
+    meta: Meta = defaultMeta
+): DirectTaskDependency<T> =
+    DirectTaskDependency(task, meta, placement).also { add(it) }
+
+fun <T : Any> TaskDependencyContainer.dependsOn(
+    task: Task<T>,
+    placement: String,
+    meta: Meta = defaultMeta
+): DirectTaskDependency<T> =
+    DirectTaskDependency(task, meta, placement.toName()).also { add(it) }
+
+fun <T : Any> TaskDependencyContainer.dependsOn(
+    task: Task<T>,
+    placement: Name = EmptyName,
+    metaBuilder: MetaBuilder.() -> Unit
+): DirectTaskDependency<T> =
+    dependsOn(task, placement, buildMeta(metaBuilder))
+
+/**
+ * Add custom data dependency
+ */
+fun TaskDependencyContainer.data(action: DataFilter.() -> Unit): DataDependency =
+    DataDependency(DataFilter.build(action)).also { add(it) }
+
+/**
+ * User-friendly way to add data dependency
+ */
+fun TaskDependencyContainer.data(pattern: String? = null, from: String? = null, to: String? = null): DataDependency =
+    data {
+        pattern?.let { this.pattern = it }
+        from?.let { this.from = it }
+        to?.let { this.to = it }
+    }
+
+/**
+ * Add all data as root node
+ */
+fun TaskDependencyContainer.allData(to: Name = EmptyName) = AllDataDependency(to).also { add(it) }
+
 /**
  * A builder for [TaskModel]
  */
-@TaskBuildScope
-class TaskModelBuilder(val name: Name, meta: Meta = EmptyMeta) {
+class TaskModelBuilder(val name: Name, meta: Meta = EmptyMeta) : TaskDependencyContainer {
     /**
      * Meta for current task. By default uses the whole input meta
      */
     var meta: MetaBuilder = meta.builder()
     val dependencies = HashSet<Dependency>()
 
+    override val defaultMeta: Meta get() = meta
+
+    override fun add(dependency: Dependency) {
+        dependencies.add(dependency)
+    }
+
     var target: String by this.meta.string(key = MODEL_TARGET_KEY, default = "")
 
-    /**
-     * Add dependency for a task defined in a workspace and resolved by
-     */
-    fun dependsOn(name: Name, meta: Meta = this.meta, placement: Name = EmptyName) {
-        dependencies.add(WorkspaceTaskDependency(name, meta, placement))
-    }
-
-    fun dependsOn(name: String, meta: Meta = this.meta, placement: Name = EmptyName) =
-        dependsOn(name.toName(), meta, placement)
-
-    fun dependsOn(task: Task<*>, meta: Meta = this.meta, placement: Name = EmptyName) {
-        dependencies.add(DirectTaskDependency(task, meta, placement))
-    }
-
-    fun dependsOn(task: Task<*>, placement: Name = EmptyName, metaBuilder: MetaBuilder.() -> Unit) =
-        dependsOn(task.name, buildMeta(metaBuilder), placement)
-
-    /**
-     * Add custom data dependency
-     */
-    fun data(action: DataFilter.() -> Unit) {
-        dependencies.add(DataDependency(DataFilter.build(action)))
-    }
-
-    /**
-     * User-friendly way to add data dependency
-     */
-    fun data(pattern: String? = null, from: String? = null, to: String? = null) = data {
-        pattern?.let { this.pattern = it }
-        from?.let { this.from = it }
-        to?.let { this.to = it }
-    }
-
-    /**
-     * Add all data as root node
-     */
-    fun allData(to: Name = EmptyName) {
-        dependencies.add(AllDataDependency(to))
-    }
 
     fun build(): TaskModel = TaskModel(name, meta.seal(), dependencies)
 }
diff --git a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/Workspace.kt b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/Workspace.kt
index f5f0f3a6..31da6c56 100644
--- a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/Workspace.kt
+++ b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/Workspace.kt
@@ -1,6 +1,8 @@
 package hep.dataforge.workspace
 
+import hep.dataforge.context.Context
 import hep.dataforge.context.ContextAware
+import hep.dataforge.context.Global
 import hep.dataforge.data.Data
 import hep.dataforge.data.DataNode
 import hep.dataforge.data.dataSequence
@@ -56,6 +58,8 @@ interface Workspace : ContextAware, Provider {
 
     companion object {
         const val TYPE = "workspace"
+        operator fun invoke(parent: Context = Global, block: SimpleWorkspaceBuilder.() -> Unit): SimpleWorkspace =
+            SimpleWorkspaceBuilder(parent).apply(block).build()
     }
 }
 
@@ -73,3 +77,6 @@ fun Workspace.run(task: String, meta: Meta) =
 
 fun Workspace.run(task: String, block: MetaBuilder.() -> Unit = {}) =
     run(task, buildMeta(block))
+
+fun <T: Any> Workspace.run(task: Task<T>, metaBuilder: MetaBuilder.() -> Unit = {}): DataNode<T> =
+    run(task, buildMeta(metaBuilder))
\ No newline at end of file
diff --git a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/WorkspaceBuilder.kt b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/WorkspaceBuilder.kt
index 6bc4900a..75a660cd 100644
--- a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/WorkspaceBuilder.kt
+++ b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/WorkspaceBuilder.kt
@@ -8,6 +8,7 @@ import hep.dataforge.meta.*
 import hep.dataforge.names.EmptyName
 import hep.dataforge.names.Name
 import hep.dataforge.names.isEmpty
+import hep.dataforge.names.toName
 import kotlin.jvm.JvmName
 import kotlin.reflect.KClass
 
@@ -70,9 +71,7 @@ fun <T : Any> WorkspaceBuilder.task(
     name: String,
     type: KClass<out T>,
     builder: TaskBuilder<T>.() -> Unit
-): Task<T> {
-    return TaskBuilder(name, type).apply(builder).build().also { tasks.add(it) }
-}
+): Task<T> = TaskBuilder(name.toName(), type).apply(builder).build().also { tasks.add(it) }
 
 inline fun <reified T : Any> WorkspaceBuilder.task(
     name: String,
diff --git a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/taskTemplates.kt b/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/taskTemplates.kt
deleted file mode 100644
index a2602da3..00000000
--- a/dataforge-workspace/src/commonMain/kotlin/hep/dataforge/workspace/taskTemplates.kt
+++ /dev/null
@@ -1,5 +0,0 @@
-package hep.dataforge.workspace
-
-//fun <T1: Any, T2: Any, R: Any> TaskBuilder.zip(
-////    val firstNo
-////) = rawTransform {  }
\ No newline at end of file
diff --git a/dataforge-workspace/src/jvmTest/kotlin/hep/dataforge/workspace/SimpleWorkspaceTest.kt b/dataforge-workspace/src/jvmTest/kotlin/hep/dataforge/workspace/SimpleWorkspaceTest.kt
index 2adc27b0..dd72a905 100644
--- a/dataforge-workspace/src/jvmTest/kotlin/hep/dataforge/workspace/SimpleWorkspaceTest.kt
+++ b/dataforge-workspace/src/jvmTest/kotlin/hep/dataforge/workspace/SimpleWorkspaceTest.kt
@@ -4,7 +4,6 @@ import hep.dataforge.context.PluginTag
 import hep.dataforge.data.*
 import hep.dataforge.meta.boolean
 import hep.dataforge.meta.get
-import hep.dataforge.names.asName
 import org.junit.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
@@ -22,7 +21,7 @@ class SimpleWorkspaceTest {
         override val tasks: Collection<Task<*>> = listOf(contextTask)
     }
 
-    val workspace = SimpleWorkspace.build {
+    val workspace = Workspace {
 
         context {
             plugin(testPlugin)
@@ -35,10 +34,7 @@ class SimpleWorkspaceTest {
         }
 
 
-        val square = task("square") {
-            model {
-                allData()
-            }
+        val square = task<Int>("square") {
             pipe<Int> { data ->
                 if (meta["testFlag"].boolean == true) {
                     println("flag")
@@ -48,24 +44,21 @@ class SimpleWorkspaceTest {
             }
         }
 
-        val linear = task("linear") {
-            model {
-                allData()
-            }
+        val linear = task<Int>("linear") {
             pipe<Int> { data ->
                 context.logger.info { "Starting linear on $data" }
                 data * 2 + 1
             }
         }
 
-        val fullSquare = task("fullsquare") {
+        val fullSquare = task<Int>("fullsquare") {
             model {
-                dependsOn("square", placement = "square".asName())
-                dependsOn("linear", placement = "linear".asName())
+                val squareDep = dependsOn(square, placement = "square")
+                val linearDep = dependsOn(linear, placement = "linear")
             }
-            transform<Any> { data ->
-                val squareNode = data["square"].filterIsInstance<Int>().node!!
-                val linearNode = data["linear"].filterIsInstance<Int>().node!!
+            transform { data ->
+                val squareNode = data["square"].node!!.cast<Int>()//squareDep()
+                val linearNode = data["linear"].node!!.cast<Int>()//linearDep()
                 return@transform DataNode(Int::class) {
                     squareNode.dataSequence().forEach { (name, _) ->
                         val newData = Data {
@@ -79,9 +72,9 @@ class SimpleWorkspaceTest {
             }
         }
 
-        task("sum") {
+        task<Int>("sum") {
             model {
-                dependsOn("square")
+                dependsOn(square)
             }
             join<Int> { data ->
                 context.logger.info { "Starting sum" }
@@ -89,10 +82,7 @@ class SimpleWorkspaceTest {
             }
         }
 
-        task<Double>("average") {
-            model {
-                allData()
-            }
+        val average = task<Double>("average") {
             joinByGroup<Int> { env ->
                 group("even", filter = { name, _ -> name.toString().toInt() % 2 == 0 }) {
                     result { data ->
@@ -111,7 +101,7 @@ class SimpleWorkspaceTest {
 
         task("delta") {
             model {
-                dependsOn("average")
+                dependsOn(average)
             }
             join<Double> { data ->
                 data["even"]!! - data["odd"]!!