heap2.nix

{ pkgs, lib }:

with pkgs.lib;
with lib;

let
  percDown =
    heap: node:
    let
      cmpLhs = heap.cmp node.val node.lhs.val;
      cmpRhs = heap.cmp node.val node.rhs.val;
      cmpLhsRhs = heap.cmp node.lhs.val node.rhs.val;
    in
    # Leaf node, we're as far down as we go
    if isLeaf node then
      node
    # rhs null, only check lhs
    # If we're greater than lhs, we should perc down one more layer. Swap vals and perc down.
    else if node.rhs == null && cmpLhs > 0 then
      node
      // {
        val = node.lhs.val;
        lhs = percDown heap (node.lhs // { val = node.val; });
      }
    # no rhs node, and we're already far enough down
    else if node.rhs == null then
      node
    # both nodes, but we're down far enough
    else if cmpLhs <= 0 && cmpRhs <= 0 then
      node
    # Otherwise, swap with the lesser of the two
    else if cmpLhsRhs < 0 then
      node
      // {
        val = node.lhs.val;
        lhs = percDown heap (node.lhs // { val = node.val; });
      }
    else
      node
      // {
        val = node.rhs.val;
        rhs = percDown heap (node.rhs // { val = node.val; });
      };

  percUp =
    size: heap: node:
    if size == 0 || size == 1 then
      node
    else
      let
        f =
          lr: node:
          let
            goRight = head lr;
            childNode = if goRight then f (tail lr) node.rhs else f (tail lr) node.lhs;
          in
          if (length lr) == 0 then
            node
          # if the child perc'd up node is less than us, perc it up one more.
          else if (heap.cmp childNode.val node.val) < 0 then
            node
            // {
              val = childNode.val;
            }
            // (
              if goRight then
                {
                  rhs = childNode // {
                    val = node.val;
                  };
                }
              else
                {
                  lhs = childNode // {
                    val = node.val;
                  };
                }
            )
          # Otherwise, we're in the right place already
          else
            node
            // (
              if goRight then
                {
                  rhs = childNode;
                }
              else
                {
                  lhs = childNode;
                }
            );
        binaryDigits = tail (toBaseDigits 2 size);
      in
      f (map (d: d == 1) binaryDigits) node;

  isLeaf = node: node.lhs == null && node.rhs == null;

  removeLastVal =
    size: node:
    let
      f =
        lr: node:
        let
          removeRhs = f (tail lr) node.rhs;
          removeLhs = f (tail lr) node.lhs;
          goRight = head lr;
        in
        if (length lr) == 0 then
          {
            val = node.val;
            node = null;
          }
        else if goRight then
          {
            node = node // {
              rhs = removeRhs.node;
            };
            val = removeRhs.val;
          }
        else
          {
            node = node // {
              lhs = removeLhs.node;
            };
            val = removeLhs.val;
          };
      # To figure out where to navigate in our complete tree, we can look at the
      # size base2 of our target size, ignoring our root.
      binaryDigits = tail (toBaseDigits 2 size);
    in
    f (map (d: d == 1) binaryDigits) node;

  addLastVal =
    size: node: val:
    let
      f =
        lr: node:
        let
          goRight = head lr;
        in
        if (length lr) == 0 then
          {
            lhs = null;
            rhs = null;
            inherit val;
          }
        else if goRight then
          node // { rhs = f (tail lr) node.rhs; }
        else
          node // { lhs = f (tail lr) node.lhs; };
      # Same as above, but our new goal size.
      binaryDigits = tail (toBaseDigits 2 (size + 1));
    in
    f (map (d: d == 1) binaryDigits) node;
in
rec {
  # mkHeap returns a min binary heap which may be used with the other functions
  # in this library.
  mkHeap = cmp: {
    inherit cmp;
    root = null;
    size = 0;
  };

  # min returns the minimum element on the heap (O(1))
  min = heap: if heap.size == 0 then throw "heap: cannot pop empty heap" else heap.root.val;

  # pop removes the minimum element
  # it returns an attr of { heap = heapWithMinRemoved; val = minVal; }
  pop =
    heap:
    if heap.size == 1 then
      {
        val = min heap;
        heap = {
          cmp = heap.cmp;
          size = 0;
          root = null;
        };
      }
    else
      let
        val = min heap;
        # Remove the last node and get its value so we can swap it into the root.
        lastVal = removeLastVal heap.size heap.root;
        # Put the last val in as the root val and reheapify
        heap' = {
          cmp = heap.cmp;
          size = heap.size - 1;
          root = lastVal.node // {
            val = lastVal.val;
          };
        };
        # Reheapify
        heap'' = heap' // {
          root = percDown heap' heap'.root;
        };
      in
      {
        inherit val;
        heap = heap'';
      };

  # insert a value
  insert =
    heap: el:
    let
      heap' = {
        cmp = heap.cmp;
        size = heap.size + 1;
        root = addLastVal heap.size heap.root el;
      };
    in
    heap' // { root = percUp heap'.size heap' heap'.root; };
}