From 0989888179e1d5b4eea7f87f67963768f35d2c32 Mon Sep 17 00:00:00 2001
From: Gabriel Arazas <foodogsquared@foodogsquared.one>
Date: Fri, 13 Sep 2024 17:07:27 +0800
Subject: [PATCH] lib/trivial: add metric and binary unit conversion functions

---
 lib/trivial.nix       | 60 ++++++++++++++++++++++++++++++++++++++++++
 tests/lib/trivial.nix | 61 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 121 insertions(+)

diff --git a/lib/trivial.nix b/lib/trivial.nix
index a8bd1651..82f0f084 100644
--- a/lib/trivial.nix
+++ b/lib/trivial.nix
@@ -1,6 +1,17 @@
 { pkgs, lib, self }:
 
 rec {
+  /* Force a numerical value to be a floating value.
+
+     Example:
+       toFloat 5698
+       => 5698.0
+
+       toFloat 9859.55
+       => 9859.55
+  */
+  toFloat = x: x / 1.0;
+
   /* Count the attributes with the given predicate.
 
      Examples:
@@ -130,4 +141,53 @@ rec {
   */
   binaryPrefixMultiplier = c:
     self.math.pow 2 (binaryPrefixExponent c);
+
+  /* Parse the given string containing the size into its appropriate value.
+     Returns the value in number of bytes.
+
+     Example:
+      parseBytesSizeIntoInt "4GiB"
+      => 4294967296
+
+      parseBytesSizeIntoInt "2 MB"
+      => 2000000
+
+      parseBytesSizeIntoInt "2 Mb"
+      => 250000
+  */
+  parseBytesSizeIntoInt = str:
+    let
+      matches = builtins.match "([[:digit:]]+)[[:space:]]*([[:alpha:]]{1})(i?[B|b])" str;
+      numeral = lib.toInt (lib.lists.head matches);
+      prefix = lib.lists.elemAt matches 1;
+      suffix = lib.lists.last matches;
+      isBinary = lib.hasPrefix "i" suffix;
+
+      multiplier =
+        let
+          multiplierFn = if isBinary then binaryPrefixMultiplier else metricPrefixMultiplier;
+        in
+          multiplierFn prefix;
+      bitDivider = if lib.hasSuffix "b" suffix then 8 else 1;
+    in
+      numeral * multiplier / bitDivider;
+
+  /*
+    Given an attrset of unit size object, return the size in bytes.
+
+    Example:
+      unitsToInt { size = 4; prefix = "G"; type = "binary"; }
+      => 4294967296
+
+      unitsToInt { size = 4; prefix = "G"; type = "metric"; }
+      => 4000000000
+  */
+  unitsToInt = { size, prefix, type ? "binary" }:
+    let
+      multiplierFn =
+        if type == "binary" then binaryPrefixMultiplier
+        else if type == "metric" then metricPrefixMultiplier
+        else builtins.throw "no multiplier type ${type}";
+    in
+      size * (multiplierFn prefix);
 }
diff --git a/tests/lib/trivial.nix b/tests/lib/trivial.nix
index c27218c7..ae63e92a 100644
--- a/tests/lib/trivial.nix
+++ b/tests/lib/trivial.nix
@@ -1,6 +1,16 @@
 { pkgs, lib, self }:
 
 lib.runTests {
+  testToFloat = {
+    expr = self.trivial.toFloat 4;
+    expected = 4.0;
+  };
+
+  testToFloat2 = {
+    expr = self.trivial.toFloat 5.5;
+    expected = 5.5;
+  };
+
   testCountAttrs = {
     expr = self.trivial.countAttrs (n: v: v?enable && v.enable) {
       hello.enable = true;
@@ -69,4 +79,55 @@ lib.runTests {
     expr = self.trivial.binaryPrefixMultiplier "G";
     expected = 1073741824;
   };
+
+  testParseBytesSizeIntoInt = {
+    expr = self.trivial.parseBytesSizeIntoInt "3GB";
+    expected = 3 * (self.trivial.metricPrefixMultiplier "G");
+  };
+
+  testParseBytesSizeIntoInt2 = {
+    expr = self.trivial.parseBytesSizeIntoInt "5MiB";
+    expected = 5 * (self.trivial.binaryPrefixMultiplier "M");
+  };
+
+  testParseBytesSizeIntoInt3 = {
+    expr = self.trivial.parseBytesSizeIntoInt "5MB";
+    expected = 5 * (self.trivial.metricPrefixMultiplier "M");
+  };
+
+  testParseBytesSizeIntoInt4 = {
+    expr = self.trivial.parseBytesSizeIntoInt "2 TiB";
+    expected = 2 * (self.trivial.binaryPrefixMultiplier "T");
+  };
+
+  testParseBytesSizeIntoInt5 = {
+    expr = self.trivial.parseBytesSizeIntoInt "2 Tib";
+    expected = 2 * (self.trivial.binaryPrefixMultiplier "T") / 8;
+  };
+
+  testUnitsToInt = {
+    expr = self.trivial.unitsToInt {
+      size = 4.5;
+      prefix = "G";
+      type = "metric";
+    };
+    expected = 4500000000;
+  };
+
+  testUnitsToInt2 = {
+    expr = self.trivial.unitsToInt {
+      size = 4.5;
+      prefix = "G";
+    };
+    expected = 4831838208;
+  };
+
+  testUnitsToInt3 = {
+    expr = self.trivial.unitsToInt {
+      size = 532;
+      prefix = "M";
+      type = "metric";
+    };
+    expected = 532000000;
+  };
 }