From 6ffba900ebaa6b462fbbae8472234f26e656e63f Mon Sep 17 00:00:00 2001
From: Gabriel Arazas <foodogsquared@foodogsquared.one>
Date: Thu, 12 Sep 2024 15:49:56 +0800
Subject: [PATCH] lib/trivial: add metric and binary prefix conversion
 functions

---
 lib/trivial.nix       | 102 +++++++++++++++++++++++++++++++++++++++++-
 tests/lib/trivial.nix |  45 +++++++++++++++++++
 2 files changed, 146 insertions(+), 1 deletion(-)

diff --git a/lib/trivial.nix b/lib/trivial.nix
index eb69b726..a8bd1651 100644
--- a/lib/trivial.nix
+++ b/lib/trivial.nix
@@ -1,6 +1,6 @@
 { pkgs, lib, self }:
 
-{
+rec {
   /* Count the attributes with the given predicate.
 
      Examples:
@@ -30,4 +30,104 @@
     })
     { ok = { }; notOk = { }; }
     attrs;
+
+  /* Convenient function for converting bits to bytes.
+
+     Example:
+       bitsToBytes 1600
+       => 200
+  */
+  bitsToBytes = x: x / 8.0;
+
+  /* Gives the exponent with the associated SI prefix.
+
+     Example:
+       SIPrefixExponent "M"
+       => 6
+
+       SIPrefixExponent "Q"
+       => 30
+  */
+  SIPrefixExponent = c:
+    let
+      prefixes = {
+        Q = 30;
+        R = 27;
+        Y = 24;
+        Z = 21;
+        E = 18;
+        P = 15;
+        T = 12;
+        G = 9;
+        M = 6;
+        k = 3;
+        h = 2;
+        da = 1;
+        d = -1;
+        c = -2;
+        m = -3;
+        "μ" = -6;
+        n = -9;
+        p = -12;
+        f = -15;
+        a = -18;
+        z = -21;
+        y = -24;
+        r = -27;
+        q = -30;
+      };
+    in
+      prefixes.${c};
+
+  /* Gives the multiplier for the metric units.
+
+     Example:
+      metricPrefixMultiplier "M"
+      => 1000000
+
+      metricPrefixMultiplier "G"
+      => 1000000000
+  */
+  metricPrefixMultiplier = c:
+    self.math.pow 10 (SIPrefixExponent c);
+
+  /* Gives the exponent with the associated binary prefix.
+
+     As an implementation detail, we don't follow the proper IEC unit prefixes
+     and instead mapping this to the SI prefix for convenience.
+
+     Example:
+       SIPrefixExponent "M"
+       => 6
+
+       SIPrefixExponent "Q"
+       => 30
+  */
+  binaryPrefixExponent = c:
+    let
+      prefixes = {
+        Y = 80;
+        Z = 70;
+        E = 60;
+        P = 50;
+        T = 40;
+        G = 30;
+        M = 20;
+        K = 10;
+      };
+    in
+      prefixes.${c};
+
+  /* Gives the multiplier for the given byte unit. Essentially returns the
+     value in number of bytes.
+
+     Example:
+       binaryPrefixMultiplier "M"
+       => 1048576
+
+       binaryPrefixMultiplier "G"
+       => 1.099511628×10¹²
+  */
+  binaryPrefixMultiplier = c:
+    self.math.pow 2 (binaryPrefixExponent c);
 }
diff --git a/tests/lib/trivial.nix b/tests/lib/trivial.nix
index ccf8bfdd..c27218c7 100644
--- a/tests/lib/trivial.nix
+++ b/tests/lib/trivial.nix
@@ -24,4 +24,49 @@ lib.runTests {
       notOk = { e = 5; f = 7; };
     };
   };
+
+  testSIPrefixExponent = {
+    expr = self.trivial.SIPrefixExponent "M";
+    expected = 6;
+  };
+
+  testSIPrefixExponent2 = {
+    expr = self.trivial.SIPrefixExponent "G";
+    expected = 9;
+  };
+
+  testMetricPrefixMultiplier = {
+    expr = self.trivial.metricPrefixMultiplier "M";
+    expected = 1000000;
+  };
+
+  testMetricPrefixMultiplier2 = {
+    expr = self.trivial.metricPrefixMultiplier "G";
+    expected = 1000000000;
+  };
+
+  testBinaryPrefixMultiplier = {
+    expr = self.trivial.binaryPrefixMultiplier "M";
+    expected = 1048576;
+  };
+
+  testBinaryPrefixExponent = {
+    expr = self.trivial.binaryPrefixExponent "M";
+    expected = 20;
+  };
+
+  testBinaryPrefixExponent2 = {
+    expr = self.trivial.binaryPrefixExponent "G";
+    expected = 30;
+  };
+
+  testBinaryPrefixMultiplier2 = {
+    expr = self.trivial.binaryPrefixMultiplier "K";
+    expected = 1024;
+  };
+
+  testBinaryPrefixMultiplier3 = {
+    expr = self.trivial.binaryPrefixMultiplier "G";
+    expected = 1073741824;
+  };
 }