heap.psl

import JJH::*
import *

interface Heapable<> is
  op "=?"(Left, Right : Heapable) -> Orders::Full_Ordering
  // Also assignable, but everything is Assignable<>
end interface Heapable;

//--------TODO--------------------------------------------------------
//  Wait for ability to return multiple things
//--------------------------------------------------------------------

// log(n) insertion and removal heap implemented on ZVector
// Left child is located at 2 * n + 1
// Right child is located at 2 * n + 2 where n is the index of the current node
// Parent is located at (n - 1) / 2
// This allows the tree to be represented as an array

concurrent interface Heap<E is Heapable<>; Pole : Orders::Polarity := #less> is
  // create empty Heap
  op "[]"() -> Heap is Create;
  // adds element to heap
  op "|="(locked var H : Heap; Elem : E) {Length(H') == Length(H) + 1};

  op "|"(Left, Right : E) -> Heap;
  op "|"(Left : Heap; Right : E) -> Heap;
  op "|"(Left : E; Right : Heap) -> Heap;
  op "|"(Left : Heap; Right : Heap) -> Heap;

  op "magnitude"(H : Heap) -> Univ_Integer is Length;

  // shallow copy. Fresh array, not fresh elements
  func Copy(locked H : Heap) -> Heap;
  // creates a new empty heap with the given polarity
  func Create() -> Heap;
  // returns number of elements
  func Length(H : Heap) -> Univ_Integer;
  // pops highest priority from heap
  // return null if H is empty
  func Pop(locked var H : Heap) -> optional E {Length(H') <= Length(H)};
end interface Heap;

concurrent class Heap<E is Heapable<>; Pole : Orders::Polarity := #less> is
    // a polarity of #less will pop smallest first, #greater will pop largest first
    var Data : ZVector<E>; //Data storage for Heap
    var Size : Univ_Integer; //Number of non-null elements in Data

    // Preserve Heap Order (Parent relation to both children must match Pole)
    // Start at bottom and recurse up the tree
    func Preserve_Order_Up(locked var H : Heap; N : Integer) is
      if N == 0 then
        return;
      end if;
      const Parent_Index := (N - 1) / 2;
      const Child_Parent : Ordering := H.Data[N] =? H.Data[Parent_Index];
      if Child_Parent == Pole then
        H.Data[N] <=> H.Data[Parent_Index]
        Preserve_Order_Up(H, Parent_Index)
      end if;
    end func Preserve_Order_Up;

    // Recurses down the tree. Swaps nodes if necessary to preserve heap order
    func Preserve_Order_Down(locked var H : Heap; N : Univ_Integer) is
      const L := 2 * N + 1;
      if L >= H.Size then // If we have no left child, then we have no children
        return;     // and there's nothing to check
      end if;
      const R := 2 * N + 2;
      const Has_Right : Boolean := R < H.Size; //Check for right child
      const Left_To_Parent : Ordering := H.Data[L] =? H.Data[N];
      var Right_To_Parent : Ordering;
      if Has_Right then
        Right_To_Parent := H.Data[R] =? H.Data[N];
      end if;
      // "and then" used in this if statement (short-circuit boolean operator)
      // so that we won't use Right_To_Parent if it's null
      if Left_To_Parent == Pole or
          (Has_Right and then Right_To_Parent == Pole) then
        if Has_Right then // One or both children out of order
          const Left_To_Right := H.Data[L] =? H.Data[R];
          const Right_To_Left := H.Data[R] =? H.Data[L];
          // swap with the higher priority (lesser if Pole == #less
          // greater if Pole == #greater)
          if Left_To_Right == Pole then
            H.Data[N] <=> H.Data[L];
            Preserve_Order_Down(H, L); // Recurse down left subtree
          elsif Right_To_Left == Pole or Right_To_Left == #equal then
            H.Data[N] <=> H.Data[R];
            Preserve_Order_Down(H, R); // Recurse down right subtree
          end if;
        else // Something is out of order and there is no right child
          // It must be the left child then
          H.Data[N] <=> H.Data[L];
          // No need to recurse because
          // Rows are filled to completion before next row
        end if;
      end if;
    end func Preserve_Order_Down;

    // Move all elements from smaller into larger
    func Move(var Smaller : Heap; var Larger : Heap) -> Heap is
      for i in 0 ..< Length(Smaller) concurrent loop
        Larger |= Smaller.Pop();
      end loop;
      return Larger;
    end func Move;

  exports
    op "|="(locked var H : Heap; Elem : E) {Length(H') == Length(H) + 1} is
      if H.Data.Length() == H.Size then //no nulls, concatenate onto end
        H.Data |= Elem;
      else // there are some nulls at the end, replace the first one
        H.Data[H.Size] := Elem;
      end if;
      H.Size += 1;
      // Traverse up the tree Preserving heap order
      Preserve_Order_Up(H, H.Size - 1);
    end op "|=";

    func Create() -> Heap is
      return (Data => [], Size => 0);
    end func Create;

    op "|"(Left, Right : E) -> Heap is
      var Result := Create();
      Result |= Left;
      Result |= Right;
      return Result;
    end op "|";
    op "|"(Left : Heap; Right : E) -> Heap is
      var Result := Copy(Left);
      Result |= Right;
      return Result;
    op "|"(Left : E; Right : Heap) -> Heap is 
      return Right | Left 
    end op "|";
    op "|"(Left : Heap; Right : Heap) -> Heap is
      if Length(Left) < Length(Right) then
        var Smaller := Copy(Left);
        var Larger := Copy(Right);
        return Move(Smaller, Larger);
      else
        var Smaller := Copy(Right);
        var Larger := Copy(Left);
        return Move(Smaller, Larger);
      end if;
    end op "|";

    func Copy(locked H : Heap) -> Heap is
      return (Data => H.Data[0 ..< Length(H.Data)], Size => H.Size);
    end func Copy;

    func Length(H : Heap) -> Univ_Integer is
      return H.Size;
    end func Length;

    func Pop(locked var H : Heap) -> optional E {Length(H') <= Length(H)} is
      if H.Size == 0 then
        return null;
      end if;
      H.Data[0] <=> H.Data[H.Size - 1];
      const To_Return <== H.Data[H.Size - 1];
      H.Size -= 1;
      Preserve_Order_Down(H, 0);
      return To_Return;
    end func Pop;
end class Heap;