feat(nix/yants): add restrict r/2050

`restrict` uses a predicate function to restrict a type, giving the
restricting a descriptive name in the process.

First, the wrapped type definition is checked (e.g. int) and then the
value is checked with the predicate, so the predicate can already
depend on the value being of the wrapped type.

Change-Id: Ic3edde45a8f34c31bc164414580d0a1aa5a821d5
Reviewed-on: https://cl.tvl.fyi/c/depot/+/2312
Tested-by: BuildkiteCI
Reviewed-by: Profpatsch <mail@profpatsch.de>
Reviewed-by: tazjin <mail@tazj.in>

2 files changed, 30 insertions, 0 deletions

diff --git a/nix/yants/default.nix b/nix/yants/default.nix
index 6da99fa3c8c4..3e9a4c00a29a 100644
--- a/nix/yants/default.nix
+++ b/nix/yants/default.nix
@@ -296,4 +296,27 @@ in lib.fix (self: {
     in sig: func: if length sig < 2
       then (throw "Signature must at least have two types (a -> b)")
       else defun' sig func;
+
+  # Restricting types
+  #
+  # `restrict` wraps a type `t`, and uses a predicate `pred` to further
+  # restrict the values, giving the restriction a descriptive `name`.
+  #
+  # First, the wrapped type definition is checked (e.g. int) and then the
+  # value is checked with the predicate, so the predicate can already
+  # depend on the value being of the wrapped type.
+  restrict = name: pred: t:
+    let restriction = "${t.name}[${name}]"; in typedef' {
+    name = restriction;
+    checkType = v:
+      let res = t.checkType v;
+      in
+        if !(t.checkToBool res)
+        then res
+        else {
+          ok = pred v;
+          err = "${prettyPrint v} does not conform to restriction '${restriction}'";
+        };
+  };
+
 })
diff --git a/nix/yants/tests/default.nix b/nix/yants/tests/default.nix
index da539ca3562b..8ad306df932b 100644
--- a/nix/yants/tests/default.nix
+++ b/nix/yants/tests/default.nix
@@ -92,4 +92,11 @@ deepSeq rec {
     (struct { a = int; b = option string; })
     (sum { a = int; b = option string; })
   ];
+
+  testRestrict = [
+    ((restrict "< 42" (i: i < 42) int) 25)
+    ((restrict "not too long" (l: builtins.length l < 3) (list int)) [ 1 2 ])
+    (list (restrict "eq 5" (v: v == 5) any) [ 5 5 5 ])
+  ];
+
 } (pkgs.writeText "yants-tests" "All tests passed!")