From 749d9ccccac3a01cd6e5eff900f6e532d526bed5 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Thu, 14 Feb 2019 11:37:47 +0100
Subject: [PATCH] autopilot/combinedattach: add SetNodeScores

This commit adds a method SetNodesScores to the WeightedCombAttachment
heuristic.

Since the heuristic keeps a list of sub-heuristics, it will attempt to
recursively apply the scores to the sub heuristics.
---
 autopilot/combinedattach.go | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/autopilot/combinedattach.go b/autopilot/combinedattach.go
index fa6a4830..f61ba33e 100644
--- a/autopilot/combinedattach.go
+++ b/autopilot/combinedattach.go
@@ -126,3 +126,34 @@ func (c *WeightedCombAttachment) NodeScores(g ChannelGraph, chans []Channel,
 
 	return scores, nil
 }
+
+// SetNodeScores is used to set the internal map from NodeIDs to scores. The
+// passed scores must be in the range [0, 1.0]. The fist parameter is the name
+// of the targeted heuristic, to allow recursively target specific
+// sub-heuristics. The returned boolean indicates whether the targeted
+// heuristic was found.
+//
+// Since this heuristic doesn't keep any internal scores, it will recursively
+// apply the scores to its sub-heuristics.
+func (c *WeightedCombAttachment) SetNodeScores(targetHeuristic string,
+	newScores map[NodeID]float64) (bool, error) {
+
+	found := false
+	for _, h := range c.heuristics {
+		// It must be ScoreSettable to be available for external
+		// scores.
+		s, ok := h.AttachmentHeuristic.(ScoreSettable)
+		if !ok {
+			continue
+		}
+
+		// Heuristic supports scoring, attempt to set them.
+		applied, err := s.SetNodeScores(targetHeuristic, newScores)
+		if err != nil {
+			return false, err
+		}
+		found = found || applied
+	}
+
+	return found, nil
+}