Nix set library

sets.nix

# A set library, using attrsets with hashed keys as internal representation.

{ lib }:

let

  /*  The empty set.

      Type: sets.empty :: <set a>

      Example:
        sets.empty
        => <set a>
  */
  empty = {};

  /*  Create a set from a list of values, based on equality.
      Duplicates are collapsed into the same value.
      Order is not preserved.

      Type: sets.fromList :: [a] -> <set a>

      Example:
        sets.fromList [ 23 false { foo = 3; } { foo = 3; } { bar = 4; } ]
        => <set | 23 false { foo = 3; } { bar = 3 }>
  */
  fromList = list: lib.foldl' (set: el: insert el set) empty list;

  /*  Create a list from a set of values.
      The elements are in a random order.

      Type: sets.toList :: <set a> -> [a]

      Example:
         sets.toList <set | 23 false { foo = 3; } { bar = 3; }>
         => [ 23 false { foo = 3; } { bar = 3; } ]
  */
  toList = lib.mapAttrsToList (_: el: el);

  /*  Insert an element into a set.

      Type: sets.insert :: a -> <set a> -> <set a>

      Example:
        sets.insert 6 (sets.insert 5 (sets.fromList [ 6 7 ]))
        => <set | 5 6 7>

  */
  insert = el: set:
    let hashed = hashElement el; in
    # this is the hash collision check; in case two values are different but map to the same key, we abort;
    # We could improve this by inserting both values into a list.
    assert (set ? ${hashed}) -> set.${hashed} == el;
    set // { ${hashed} = el; };

  /*  Check whether an element is in the set.

      Type: sets.elem :: a -> <set a> -> bool

      Example:
        sets.elem 2 (sets.fromList [ 3 4 ])
        => false
        sets.elem 3 (sets.fromList [ 3 4 ])
        => true
  */
  elem = el: set: set ? ${hashElement el};

  /*  The union of all elements in both sets.
      That is, if an element is in `left` OR `right` it is included in the resulting set.

      Type: sets.union :: <left = set a> -> <right = set a> -> <set a>

      Example:
        sets.union
          (sets.fromList [ [ "foo" ] 23 { bar.baz = {}; } ])
          (sets.fromList [ [ "bar" ] 42 { bar.baz = {}; } ])
        => <set | [ "foo" ] [ "bar" ] 23 42 { bar.baz = {}; }>
  */
  union = left: right:
    # this is right-leaning merge; if a key exists in both attrsets, the value is by definition equal
    left // right;

  /*  The intersection of the elements in both sets.
      That is, if an element is in `left` AND `right` it is included in the resulting set.

      Type: sets.intersection :: <left = set a> -> <right = set a> -> <set a>

      Example:
        sets.intersection
          (sets.fromList [ [ "foo" ] 23 { bar.baz = {}; } ])
          (sets.fromList [ [ "bar" ] 42 { bar.baz = {}; } ])
        => <set | { bar.baz = {}; }>
  */
  intersection = left: right:
    filter (el: elem el right) left;

  /*  The difference of two sets.
      That is, if an element is in `set` BUT NOT in `remove` it *will* be included in the resulting set.
      If an element is in `remove` BUT NOT in `set`, it *won’t* be included in the resulting set.

      The *left* argument is `remove`, so `a - b` is `sets.difference b a`.
      This is done to keep symmetry with the other functions (the rightmost set is the “working” set).

      Type: sets.difference :: <remove = set a> -> <set = set a> -> <set a>

      Example:
        sets.difference
          (sets.fromList [ [ "foo" ] 23 { bar.baz = {}; } ])
          (sets.fromList [ [ "bar" ] 42 { bar.baz = {}; } ])
        => <set | [ "bar" ] 42>
  */
  difference = remove: set:
    filter (el: !(elem el remove)) set;

  /*  Filter a set based on a predicate.
      Elements for which the predicate returns `true` are kept.

      Type: sets.filter :: (a -> Bool) -> <set a> -> <set a>

      Example:
        sets.filter
          (x: x > 42)
          (sets.fromList [ 1 20 42 57 ])
        => <set | 57>
  */
  filter = pred: set: lib.filterAttrs (_: v: pred v) set;


  ## INTERNAL

  # we use md5 as a good compromise between result length and collision probability;
  # since this application is not security critical, we don’t care that results can be forged.
  # We check for collisions before inserting elements into the set.
  keyHashFunction = "md5";

  # hash a value in order to add it as an element to the set.
  hashElement = el:
    # elToString *must* be injective, otherwise different values map to the same string and hash.
    # This is achieved by giving each type a different prefix char.
    let
      elToString = v:
             if builtins.isInt v then "i${toString v}"
        else if builtins.isFloat v then "f${toString v}"
        else if builtins.isString v then "s${v}"
        else if true == v then "b1"
        else if false == v then "b0"
        else if null == v then "n"
        else if builtins.isList  v then "l${lib.concatMapStrings elToString v}"
        else if builtins.isAttrs v then "a${lib.concatStrings (lib.mapAttrsToList (key: val: "k${key}v${elToString val}") v)}"
        else if builtins.isFunction v then abort "sets.hashElement: cannot hash a function (${lib.generators.toPretty {} v})"
        else abort "sets.hashElement: don’t know how to hash ${lib.generators.toPretty {} v}";

    in builtins.hashString keyHashFunction (elToString el);

in {
  inherit
    empty
    fromList
    toList
    insert
    elem
    union
    intersection
    difference
    filter
    ;
}