From 3ea86bfbf430c70194d7fc4132139f464522fe0f Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Mon, 9 Jun 2025 12:03:42 +0200
Subject: [PATCH] graph/db: let GraphCache.AddChannel take cached edge polies

So that the call-sites of this method dont necessarily need to fetch all
the data required to populate the full ChannelEdgePolicy.
---
 graph/db/graph.go            | 12 ++++++++++--
 graph/db/graph_cache.go      | 14 +++++---------
 graph/db/graph_cache_test.go | 23 ++++++++++++-----------
 3 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/graph/db/graph.go b/graph/db/graph.go
index 30486adae..b414ff696 100644
--- a/graph/db/graph.go
+++ b/graph/db/graph.go
@@ -177,7 +177,11 @@ func (c *ChannelGraph) populateCache() error {
 	err = c.V1Store.ForEachChannel(func(info *models.ChannelEdgeInfo,
 		policy1, policy2 *models.ChannelEdgePolicy) error {
 
-		c.graphCache.AddChannel(info, policy1, policy2)
+		c.graphCache.AddChannel(
+			info,
+			models.NewCachedPolicy(policy1),
+			models.NewCachedPolicy(policy2),
+		)
 
 		return nil
 	})
@@ -347,7 +351,11 @@ func (c *ChannelGraph) MarkEdgeLive(chanID uint64) error {
 
 		info := infos[0]
 
-		c.graphCache.AddChannel(info.Info, info.Policy1, info.Policy2)
+		c.graphCache.AddChannel(
+			info.Info,
+			models.NewCachedPolicy(info.Policy1),
+			models.NewCachedPolicy(info.Policy2),
+		)
 	}
 
 	return nil
diff --git a/graph/db/graph_cache.go b/graph/db/graph_cache.go
index 43069dd46..f07c0d38b 100644
--- a/graph/db/graph_cache.go
+++ b/graph/db/graph_cache.go
@@ -115,7 +115,7 @@ func (c *GraphCache) AddNodeFeatures(node route.Vertex,
 // and policy flags automatically. The policy will be set as the outgoing policy
 // on one node and the incoming policy on the peer's side.
 func (c *GraphCache) AddChannel(info *models.ChannelEdgeInfo,
-	policy1 *models.ChannelEdgePolicy, policy2 *models.ChannelEdgePolicy) {
+	policy1, policy2 *models.CachedEdgePolicy) {
 
 	if info == nil {
 		return
@@ -147,21 +147,17 @@ func (c *GraphCache) AddChannel(info *models.ChannelEdgeInfo,
 	// of node 2 then we have the policy 1 as seen from node 1.
 	if policy1 != nil {
 		fromNode, toNode := info.NodeKey1Bytes, info.NodeKey2Bytes
-		if policy1.ToNode != info.NodeKey2Bytes {
+		if !policy1.IsNode1() {
 			fromNode, toNode = toNode, fromNode
 		}
-		c.UpdatePolicy(
-			models.NewCachedPolicy(policy1), fromNode, toNode,
-		)
+		c.UpdatePolicy(policy1, fromNode, toNode)
 	}
 	if policy2 != nil {
 		fromNode, toNode := info.NodeKey2Bytes, info.NodeKey1Bytes
-		if policy2.ToNode != info.NodeKey1Bytes {
+		if policy2.IsNode1() {
 			fromNode, toNode = toNode, fromNode
 		}
-		c.UpdatePolicy(
-			models.NewCachedPolicy(policy2), fromNode, toNode,
-		)
+		c.UpdatePolicy(policy2, fromNode, toNode)
 	}
 }
 
diff --git a/graph/db/graph_cache_test.go b/graph/db/graph_cache_test.go
index 96318366d..d9ad1dad8 100644
--- a/graph/db/graph_cache_test.go
+++ b/graph/db/graph_cache_test.go
@@ -43,20 +43,21 @@ func TestGraphCacheAddNode(t *testing.T) {
 			FeeRate: 20,
 		}
 
-		outPolicy1 := &models.ChannelEdgePolicy{
+		outPolicy1 := &models.CachedEdgePolicy{
 			ChannelID:    1000,
 			ChannelFlags: lnwire.ChanUpdateChanFlags(channelFlagA),
-			ToNode:       nodeB,
+			ToNodePubKey: func() route.Vertex {
+				return nodeB
+			},
 			// Define an inbound fee.
 			InboundFee: fn.Some(inboundFee),
-			ExtraOpaqueData: []byte{
-				253, 217, 3, 8, 0, 0, 0, 10, 0, 0, 0, 20,
-			},
 		}
-		inPolicy1 := &models.ChannelEdgePolicy{
+		inPolicy1 := &models.CachedEdgePolicy{
 			ChannelID:    1000,
 			ChannelFlags: lnwire.ChanUpdateChanFlags(channelFlagB),
-			ToNode:       nodeA,
+			ToNodePubKey: func() route.Vertex {
+				return nodeA
+			},
 		}
 		cache := NewGraphCache(10)
 		cache.AddNodeFeatures(nodeA, lnwire.EmptyFeatureVector())
@@ -120,7 +121,7 @@ func TestGraphCacheAddNode(t *testing.T) {
 	runTest(pubKey2, pubKey1)
 }
 
-func assertCachedPolicyEqual(t *testing.T, original *models.ChannelEdgePolicy,
+func assertCachedPolicyEqual(t *testing.T, original,
 	cached *models.CachedEdgePolicy) {
 
 	require.Equal(t, original.ChannelID, cached.ChannelID)
@@ -134,7 +135,7 @@ func assertCachedPolicyEqual(t *testing.T, original *models.ChannelEdgePolicy,
 		t, original.FeeProportionalMillionths,
 		cached.FeeProportionalMillionths,
 	)
-	require.Equal(
-		t, route.Vertex(original.ToNode), cached.ToNodePubKey(),
-	)
+	if original.ToNodePubKey != nil {
+		require.Equal(t, original.ToNodePubKey(), cached.ToNodePubKey())
+	}
 }