From 0de9c0751bbf4738e9d45bc85dc714375fef4f3d Mon Sep 17 00:00:00 2001
From: oddlama <oddlama@oddlama.org>
Date: Thu, 21 Nov 2024 19:51:16 +0100
Subject: [PATCH] perf: allow caching wireguard network information to
 significantly speed up evaluation

---
 lib/wireguard.nix              | 24 ++++++++++++++++++++++--
 modules/topology-wireguard.nix |  4 ++--
 modules/wireguard.nix          |  7 +++----
 3 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/lib/wireguard.nix b/lib/wireguard.nix
index 5f68522..458f64d 100644
--- a/lib/wireguard.nix
+++ b/lib/wireguard.nix
@@ -16,6 +16,7 @@ inputs: final: prev: let
     nameValuePair
     partition
     removeSuffix
+    warn
     ;
 
   inherit
@@ -32,8 +33,8 @@ inputs: final: prev: let
 in {
   lib =
     prev.lib
-    // {
-      wireguard = userInputs: wgName: let
+    // rec {
+      wireguard.evaluateNetwork = userInputs: wgName: let
         inherit (userInputs.self) nodes;
         # Returns the given node's wireguard configuration of this network
         wgCfgOf = node: nodes.${node}.config.wireguard.${wgName};
@@ -220,5 +221,24 @@ in {
           wgQuickConfigScript
           ;
       };
+
+      wireguard.createEvalCache = userInputs: wgNames:
+        genAttrs wgNames (wireguard.evaluateNetwork userInputs);
+
+      wireguard.getNetwork = userInputs: wgName:
+        userInputs.self.wireguardEvalCache.${wgName}
+        or (
+          warn ''
+            The calculated information for the wireguard network "${wgName}" is not cached!
+            This will siginificantly increase evaluation times. Please consider pre-evaluating
+            this information by exposing it in your flake:
+
+              wireguardEvalCache.${wgName} = lib.wireguard.createEvalCache inputs [
+                "${wgName}"
+                # all other networks
+              ];
+
+          '' (wireguard.evaluateNetwork userInputs wgName)
+        );
     };
 }
diff --git a/modules/topology-wireguard.nix b/modules/topology-wireguard.nix
index 09840a0..31cc5a8 100644
--- a/modules/topology-wireguard.nix
+++ b/modules/topology-wireguard.nix
@@ -30,7 +30,7 @@ in {
     topology.networks = mkMerge (
       flip mapAttrsToList config.wireguard (
         wgName: _: let
-          inherit (lib.wireguard inputs wgName) networkCidrs;
+          inherit (lib.wireguard.getNetwork inputs wgName) networkCidrs;
         in {
           ${networkId wgName} = {
             name = mkDefault "Wireguard network '${wgName}'";
@@ -47,7 +47,7 @@ in {
       flip mapAttrsToList config.wireguard (
         wgName: wgCfg: let
           inherit
-            (lib.wireguard inputs wgName)
+            (lib.wireguard.getNetwork inputs wgName)
             participatingServerNodes
             wgCfgOf
             ;
diff --git a/modules/wireguard.nix b/modules/wireguard.nix
index 94d68a6..a3ca646 100644
--- a/modules/wireguard.nix
+++ b/modules/wireguard.nix
@@ -28,7 +28,6 @@
     optionals
     stringLength
     types
-    wireguard
     ;
 
   cfg = config.wireguard;
@@ -36,7 +35,7 @@
 
   configForNetwork = wgName: wgCfg: let
     inherit
-      (wireguard inputs wgName)
+      (lib.wireguard.getNetwork inputs wgName)
       externalPeerName
       externalPeerNamesRaw
       networkCidrs
@@ -385,7 +384,7 @@ in {
 
         ipv4 = mkOption {
           type = types.lazyOf types.net.ipv4;
-          default = types.lazyValue (wireguard inputs name).assignedIpv4Addresses.${nodeName};
+          default = types.lazyValue (lib.wireguard.getNetwork inputs name).assignedIpv4Addresses.${nodeName};
           description = ''
             The ipv4 address for this machine. If you do not set this explicitly,
             a semi-stable ipv4 address will be derived automatically based on the
@@ -397,7 +396,7 @@ in {
 
         ipv6 = mkOption {
           type = types.lazyOf types.net.ipv6;
-          default = types.lazyValue (wireguard inputs name).assignedIpv6Addresses.${nodeName};
+          default = types.lazyValue (lib.wireguard.getNetwork inputs name).assignedIpv6Addresses.${nodeName};
           description = ''
             The ipv6 address for this machine. If you do not set this explicitly,
             a semi-stable ipv6 address will be derived automatically based on the