From 013a0a1357c446d0a46b4bbd8f68512fd9223257 Mon Sep 17 00:00:00 2001
From: adisbladis <adisbladis@gmail.com>
Date: Fri, 24 Nov 2023 16:10:13 +1300
Subject: [PATCH] lib.attrsets.matchAttrs: Avoid some list allocations when
 walking structure

Benchmarks (`nix-instantiate ./. -A python3`):

- Before:
``` json
{
  "cpuTime": 0.29049500823020935,
  "envs": {
    "bytes": 4484216,
    "elements": 221443,
    "number": 169542
  },
  "gc": {
    "heapSize": 402915328,
    "totalBytes": 53086800
  },
  "list": {
    "bytes": 749424,
    "concats": 4242,
    "elements": 93678
  },
  "nrAvoided": 253991,
  "nrFunctionCalls": 149848,
  "nrLookups": 49612,
  "nrOpUpdateValuesCopied": 1587837,
  "nrOpUpdates": 10104,
  "nrPrimOpCalls": 130356,
  "nrThunks": 358981,
  "sets": {
    "bytes": 30423600,
    "elements": 1859999,
    "number": 41476
  },
  "sizes": {
    "Attr": 16,
    "Bindings": 16,
    "Env": 16,
    "Value": 24
  },
  "symbols": {
    "bytes": 236145,
    "number": 24453
  },
  "values": {
    "bytes": 10502520,
    "number": 437605
  }
}
```

- After:
``` json
{
  "cpuTime": 0.2946169972419739,
  "envs": {
    "bytes": 3315224,
    "elements": 172735,
    "number": 120834
  },
  "gc": {
    "heapSize": 402915328,
    "totalBytes": 48718432
  },
  "list": {
    "bytes": 347568,
    "concats": 4242,
    "elements": 43446
  },
  "nrAvoided": 173252,
  "nrFunctionCalls": 101140,
  "nrLookups": 73595,
  "nrOpUpdateValuesCopied": 1587837,
  "nrOpUpdates": 10104,
  "nrPrimOpCalls": 83067,
  "nrThunks": 304216,
  "sets": {
    "bytes": 29704096,
    "elements": 1831673,
    "number": 24833
  },
  "sizes": {
    "Attr": 16,
    "Bindings": 16,
    "Env": 16,
    "Value": 24
  },
  "symbols": {
    "bytes": 236145,
    "number": 24453
  },
  "values": {
    "bytes": 8961552,
    "number": 373398
  }
}
```
---
 lib/attrsets.nix   | 29 ++++++++++++++++++++---------
 lib/tests/misc.nix | 20 ++++++++++++++++++++
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/lib/attrsets.nix b/lib/attrsets.nix
index bf6c90bf1be6..14ce9c257731 100644
--- a/lib/attrsets.nix
+++ b/lib/attrsets.nix
@@ -883,7 +883,10 @@ rec {
     recursiveUpdateUntil (path: lhs: rhs: !(isAttrs lhs && isAttrs rhs)) lhs rhs;
 
 
-  /* Returns true if the pattern is contained in the set. False otherwise.
+  /*
+    Recurse into every attribute set of the first argument and check that:
+    - Each attribute path also exists in the second argument.
+    - If the attribute's value is not a nested attribute set, it must have the same value in the right argument.
 
      Example:
        matchAttrs { cpu = {}; } { cpu = { bits = 64; }; }
@@ -895,16 +898,24 @@ rec {
   matchAttrs =
     # Attribute set structure to match
     pattern:
-    # Attribute set to find patterns in
+    # Attribute set to check
     attrs:
     assert isAttrs pattern;
-    all id (attrValues (zipAttrsWithNames (attrNames pattern) (n: values:
-      let pat = head values; val = elemAt values 1; in
-      if length values == 1 then false
-      else if isAttrs pat then isAttrs val && matchAttrs pat val
-      else pat == val
-    ) [pattern attrs]));
-
+    all
+    ( # Compare equality between `pattern` & `attrs`.
+      attr:
+      # Missing attr, not equal.
+      attrs ? ${attr} && (
+        let
+          lhs = pattern.${attr};
+          rhs = attrs.${attr};
+        in
+        # If attrset check recursively
+        if isAttrs lhs then isAttrs rhs && matchAttrs lhs rhs
+        else lhs == rhs
+      )
+    )
+    (attrNames pattern);
 
   /* Override only the attributes that are already present in the old set
     useful for deep-overriding.
diff --git a/lib/tests/misc.nix b/lib/tests/misc.nix
index 06cb5e763e2c..57c1be4a2073 100644
--- a/lib/tests/misc.nix
+++ b/lib/tests/misc.nix
@@ -831,6 +831,26 @@ runTests {
     };
   };
 
+  testMatchAttrsMatchingExact = {
+    expr = matchAttrs { cpu = { bits = 64; }; } { cpu = { bits = 64; }; };
+    expected = true;
+  };
+
+  testMatchAttrsMismatch = {
+    expr = matchAttrs { cpu = { bits = 128; }; } { cpu = { bits = 64; }; };
+    expected = false;
+  };
+
+  testMatchAttrsMatchingImplicit = {
+    expr = matchAttrs { cpu = { }; } { cpu = { bits = 64; }; };
+    expected = true;
+  };
+
+  testMatchAttrsMissingAttrs = {
+    expr = matchAttrs { cpu = {}; } { };
+    expected = false;
+  };
+
   testOverrideExistingEmpty = {
     expr = overrideExisting {} { a = 1; };
     expected = {};