Use computation instead of RelationshipWeightApplier

Author: DarthMax

pregel/src/main/java/org/neo4j/gds/beta/pregel/ForkJoinComputer.java

@@ -103,7 +103,7 @@ void release() {
         Supplier<ComputeContext<CONFIG>> computeContext = () -> new ComputeContext<>(
             graph.concurrentCopy(),
             config,
-            ((PregelComputation<CONFIG>) computation)::applyRelationshipWeight,
+            computation,
             nodeValues,
             messenger,
             voteBits,
@@ -144,7 +144,7 @@ void release() {
         Supplier<BidirectionalComputeContext<CONFIG>> computeContext = () -> new BidirectionalComputeContext<>(
             graph.concurrentCopy(),
             config,
-            ((BidirectionalPregelComputation<CONFIG>) computation)::applyRelationshipWeight,
+            computation,
             nodeValues,
             messenger,
             voteBits,

pregel/src/main/java/org/neo4j/gds/beta/pregel/PartitionedComputer.java

@@ -149,7 +149,7 @@ void release() {
         var computeContext = new ComputeContext<>(
             graph,
             config,
-            ((PregelComputation<CONFIG>) computation)::applyRelationshipWeight,
+            computation,
             nodeValues,
             messenger,
             voteBits,
@@ -192,7 +192,7 @@ void release() {
         var computeContext = new BidirectionalComputeContext<>(
             graph,
             config,
-            ((BidirectionalPregelComputation<CONFIG>) computation)::applyRelationshipWeight,
+            computation,
             nodeValues,
             messenger,
             voteBits,

pregel/src/main/java/org/neo4j/gds/beta/pregel/context/ComputeContext.java

@@ -21,6 +21,7 @@
 
 import org.apache.commons.lang3.mutable.MutableInt;
 import org.neo4j.gds.api.Graph;
+import org.neo4j.gds.beta.pregel.BasePregelComputation;
 import org.neo4j.gds.beta.pregel.Messenger;
 import org.neo4j.gds.beta.pregel.NodeValue;
 import org.neo4j.gds.beta.pregel.PregelConfig;
@@ -36,24 +37,24 @@
  */
 public class ComputeContext<CONFIG extends PregelConfig> extends NodeCentricContext<CONFIG> {
 
-    final RelationshipWeightApplier relationshipWeightApplier;
     private final HugeAtomicBitSet voteBits;
 
     private final Messenger<?> messenger;
     private final MutableInt iteration;
     private final AtomicBoolean hasSendMessage;
+    protected BasePregelComputation<CONFIG> computation;
 
     public ComputeContext(Graph graph,
                           CONFIG config,
-                          RelationshipWeightApplier relationshipWeightApplier,
+                          BasePregelComputation<CONFIG> computation,
                           NodeValue nodeValue,
                           Messenger<?> messenger,
                           HugeAtomicBitSet voteBits,
                           MutableInt iteration,
                           AtomicBoolean hasSendMessage,
                           ProgressTracker progressTracker) {
         super(graph, config, nodeValue, progressTracker);
-        this.relationshipWeightApplier = relationshipWeightApplier;
+        this.computation = computation;
         this.sendMessagesFunction = config.hasRelationshipWeightProperty()
             ? this::sendToNeighborsWeighted
             : this::sendToNeighbors;
@@ -163,7 +164,7 @@ private void sendToNeighbors(long sourceNodeId, double message) {
 
     private void sendToNeighborsWeighted(long sourceNodeId, double message) {
         graph.forEachRelationship(sourceNodeId, 1.0, (ignored, targetNodeId, weight) -> {
-            sendTo(targetNodeId, relationshipWeightApplier.applyRelationshipWeight(message, weight));
+            sendTo(targetNodeId, computation.applyRelationshipWeight(message, weight));
             return true;
         });
     }
@@ -173,19 +174,14 @@ interface SendMessagesFunction {
         void sendToNeighbors(long sourceNodeId, double message);
     }
 
-    @FunctionalInterface
-    public interface RelationshipWeightApplier {
-        double applyRelationshipWeight(double nodeValue, double relationshipWeight);
-    }
-
     public static final class BidirectionalComputeContext<CONFIG extends PregelConfig> extends ComputeContext<CONFIG> implements BidirectionalNodeCentricContext {
 
         private final SendMessagesIncomingFunction sendMessagesIncomingFunction;
 
         public BidirectionalComputeContext(
             Graph graph,
             CONFIG config,
-            RelationshipWeightApplier relationshipWeightApplier,
+            BasePregelComputation<CONFIG> computation,
             NodeValue nodeValue,
             Messenger<?> messenger,
             HugeAtomicBitSet voteBits,
@@ -196,7 +192,7 @@ public BidirectionalComputeContext(
             super(
                 graph,
                 config,
-                relationshipWeightApplier,
+                computation,
                 nodeValue,
                 messenger,
                 voteBits,
@@ -226,7 +222,7 @@ private void sendToIncomingNeighbors(long sourceNodeId, double message) {
 
         private void sendToIncomingNeighborsWeighted(long sourceNodeId, double message) {
             graph.forEachInverseRelationship(sourceNodeId, 1.0, (ignored, targetNodeId, weight) -> {
-                sendTo(targetNodeId, relationshipWeightApplier.applyRelationshipWeight(message, weight));
+                sendTo(targetNodeId, computation.applyRelationshipWeight(message, weight));
                 return true;
             });
         }