Implementation of the Claim Deducer that tracks influence

java/arcs/core/analysis/DependencyNode.kt

@@ -21,10 +21,10 @@ private typealias Path = List<Identifier>
 
 /**
  * [DependencyNode]s make up a directed-acyclic-graph that describes how input handle connections
- * map to output connections in particle specs with Paxel [Expression]s.
+ * map to output connections in particle specs with [Expression]s.
  *
  * - [DependencyNode.Input] represents an input handle connection and access path.
- * - [DependencyNode.DerivedFrom] indicates that an input has been modified in the Paxel expression.
+ * - [DependencyNode.Derived] indicates that an input has been modified in the expression.
  * - [DependencyNode.AssociationNode] connects fields to other nodes in the graph. These are used to
  *   form left-hand-side / right-hand-side relations between handle connections.
  *
@@ -52,7 +52,7 @@ private typealias Path = List<Identifier>
  *   DependencyNode.AssociationNode(
  *    "ratio" to DependencyNode.AssociationNode(
  *      "trained" to DependencyNode.Input("input", "cats"),
- *      "total" to DependencyNode.DerivedFrom(
+ *      "total" to DependencyNode.Derived(
  *        DependencyNode.Input("input", "cats"),
  *        DependencyNode.Input("input", "dogs")
  *      )
@@ -84,56 +84,143 @@ private typealias Path = List<Identifier>
  */
 sealed class DependencyNode {
 
-  /** An unmodified input (from a handle connection) used in a Paxel [Expression]. */
-  data class Input(val path: Path = emptyList()) : DependencyNode() {
-    constructor(vararg fields: Identifier) : this(listOf(*fields))
-  }
+  /** Path of [Identifier]s representing access to a handle connection. */
+  abstract val path: Path
 
-  /** Represents derivation from a group of [Input]s in an [Expression]. */
-  data class DerivedFrom(val inputs: Set<Input> = emptySet()) : DependencyNode() {
-
-    constructor(vararg paths: Path) : this(paths.map { Input(it) }.toSet())
-
-    /** Produce a new [DerivedFrom] with a flattened set of [Input]s. */
-    constructor(vararg nodes: DependencyNode) : this(flatten(*nodes))
-
-    companion object {
-      /** Flatten nested sets of [DependencyNode]s.*/
-      private fun flatten(vararg nodes: DependencyNode): Set<Input> {
-        return nodes.flatMap { node ->
-          when (node) {
-            is Input -> setOf(node)
-            is DerivedFrom -> node.inputs
-            else -> throw IllegalArgumentException(
-              "Nodes must be a 'Input' or 'DerivedFrom'."
-            )
-          }
-        }.toSet()
-      }
-    }
-  }
+  /** Set of [DependencyNode]s that the current node depends on. */
+  abstract val dependency: Set<DependencyNode>
+
+  /** Set of [DependencyNode]s that bear influence on the current node. */
+  abstract val influencedBy: Set<DependencyNode>
+
+  /** Expresses influence on to the current [DependencyNode]. */
+  abstract fun influence(influenceBy: Set<DependencyNode>): DependencyNode
 
   /** Associates [Identifier]s with [DependencyNode]s. */
-  data class AssociationNode(
-    val associations: Map<Identifier, DependencyNode> = emptyMap()
+  open class AssociationNode(
+    override val path: Path = emptyList(),
+    override val dependency: Set<DependencyNode> = emptySet(),
+    override val influencedBy: Set<DependencyNode> = emptySet(),
+    val associations: Map<String, DependencyNode>
   ) : DependencyNode() {
 
-    /** Construct an [AssociationNode] from associations of [Identifier]s to [DependencyNode]s. */
-    constructor(vararg pairs: Pair<Identifier, DependencyNode>) : this(pairs.toMap())
+    /** Constructor to express associations. */
+    constructor(
+      vararg mappings: Pair<Identifier, DependencyNode>,
+      path: Path = emptyList(),
+      dependency: Set<DependencyNode> = emptySet(),
+      influencedBy: Set<DependencyNode> = emptySet()
+    ) : this(path, dependency, influencedBy, mappings.toMap())
+
+    /** Expresses influence on to the current [DependencyNode]. */
+    override fun influence(influenceBy: Set<DependencyNode>): DependencyNode = AssociationNode(
+      path,
+      dependency,
+      this.influencedBy + influenceBy,
+      associations.mapValues { (_, node) -> node.influence(influencedBy) }
+    )
 
     /** Replace the associations of an [AssociationNode] with new mappings. */
-    fun add(vararg other: Pair<Identifier, DependencyNode>): DependencyNode = AssociationNode(
+    fun add(vararg other: Pair<Identifier, DependencyNode>): AssociationNode = AssociationNode(
+      path,
+      dependency,
+      influencedBy,
       associations + other
     )
 
     /** Returns the [DependencyNode] associated with the input [Identifier]. */
-    fun lookup(key: Identifier): DependencyNode = requireNotNull(associations[key]) {
-      "Identifier '$key' is not found in AssociationNode."
+    operator fun get(key: Identifier): DependencyNode? = associations[key]
+
+    override fun equals(other: Any?): Boolean {
+      if (this === other) return true
+      if (other !is AssociationNode) return false
+
+      if (path != other.path) return false
+      if (dependency != other.dependency) return false
+      if (influencedBy != other.influencedBy) return false
+      if (associations != other.associations) return false
+
+      return true
+    }
+
+    override fun hashCode(): Int {
+      var result = path.hashCode()
+      result = 31 * result + dependency.hashCode()
+      result = 31 * result + influencedBy.hashCode()
+      result = 31 * result + associations.hashCode()
+      return result
+    }
+
+    override fun toString(): String {
+      return "AssociationNode(path=$path, dependency=$dependency, influencedBy=$influencedBy," +
+        " associations=$associations)"
+    }
+  }
+
+  /** An unmodified input (from a handle connection) used in an [Expression]. */
+  class Input(
+    override val path: Path,
+    override val dependency: Set<DependencyNode> = emptySet(),
+    override val influencedBy: Set<DependencyNode> = emptySet()
+  ) : AssociationNode(path, dependency, influencedBy, emptyMap()) {
+
+    /** Constructor to build input [Path]s. */
+    constructor(
+      vararg identifier: Identifier,
+      dependency: Set<DependencyNode> = emptySet(),
+      influencedBy: Set<DependencyNode> = emptySet()
+    ) : this(listOf(*identifier), dependency, influencedBy)
+
+    override fun toString(): String {
+      return "Input(path=$path, dependency=$dependency, influencedBy=$influencedBy)"
     }
   }
 
+  /** Represents derivation from a group of [Input]s in an [Expression]. */
+  data class Derived(
+    override val path: Path,
+    override val dependency: Set<DependencyNode> = emptySet(),
+    override val influencedBy: Set<DependencyNode> = emptySet(),
+    val inputs: Set<DependencyNode>
+  ) : DependencyNode() {
+
+    /** Constructor for Literals. */
+    constructor() : this(emptyList(), emptySet(), emptySet(), emptySet())
+
+    /** Constructor to express derivation of [Input]s. */
+    constructor(
+      vararg inputs: DependencyNode,
+      path: Path = emptyList(),
+      dependency: Set<DependencyNode> = emptySet(),
+      influencedBy: Set<DependencyNode> = emptySet()
+    ) : this(
+      path,
+      dependency,
+      influencedBy,
+      setOf(*inputs).flatten()
+    )
+
+    /** Expresses influence on to the current [DependencyNode]. */
+    override fun influence(influenceBy: Set<DependencyNode>): DependencyNode = Derived(
+      path,
+      dependency,
+      this.influencedBy + influenceBy,
+      inputs.map { node -> node.influence(influencedBy) }.toSet()
+    )
+  }
+
   companion object {
     /** A [DependencyNode] case to represent literals. */
-    val LITERAL = DerivedFrom()
+    val LITERAL = Derived()
   }
 }
+
+/** Flatten a collection of nested [DependencyNode]s. */
+fun Collection<DependencyNode>.flatten(): Set<DependencyNode> {
+  return this.flatMap { node ->
+    when (node) {
+      is DependencyNode.AssociationNode -> setOf(node) + node.associations.values.flatten()
+      is DependencyNode.Derived -> node.inputs
+    }
+  }.toSet()
+}