/* Filename : LockFreeHashSet.cvl
   Author   : YYY
   Created  : 2024-12-12
   Modified : 2025-01-17

   CIVL model of LockFreeHashSet class from "The Art of Multiprocessor
   Programming" 2nd ed, by Herlihy, Luchangco, Shavit, and Spear,
   Sec. 13.3.3 and companion code
   ch13/src/hash/LockFreeHashSet.java. Also models the BucketList
   class, Sec. 13.3.2 and companion code
   ch13/src/hash/BucketList.java.

   some Lab
   some division
   some institution
 */
#include "AtomicInteger.h"
#include "AtomicMarkableReference.cvh"
#include "Set.h"
#include "hash.cvh"
#include "tid.h"
#include "types.h"
#include <limits.h>
#include <stdbool.h>
#include <stdlib.h>

#define WORD_SIZE 3
#define LO_MASK 1
#define HI_MASK 4
#define MASK 7
#define THRESHOLD 2
#define CAPACITY 4

// =============== Start Node class ===============

typedef struct Node {
  int key;
  T value;
  AtomicMarkableReference next;
} * Node;

Node Node_create(int key, T object) {
  Node this = malloc(sizeof(struct Node));
  this->key = key;
  this->value = object;
  this->next = AtomicMarkableReference_create(NULL, false);
  return this;
}

Node Node_create_sentinel(int key) {
  return Node_create(key, 0);
}

void Node_destroy(Node n) {
  AtomicMarkableReference_destroy(n->next);
  free(n);
}

Node Node_getNext(Node this) {
  bool cMarked[1] = {false};
  bool sMarked[1] = {false};
  Node entry = AtomicMarkableReference_get(this->next, cMarked);
  while (cMarked[0]) {
    Node succ = AtomicMarkableReference_get(entry->next, sMarked);
    AtomicMarkableReference_compareAndSet(this->next, entry, succ,
                                          true, sMarked[0]);
    entry = AtomicMarkableReference_get(this->next, cMarked);
  }
  return entry;
}

// =============== End Node class ===============

// =============== Start Window class ===============

typedef struct Window {
  Node pred;
  Node curr;
} * Window;

Window Window_create(Node pred, Node curr) {
  Window this = malloc(sizeof(struct Window));
  this->pred = pred;
  this->curr = curr;
  return this;
}

void Window_destroy(Window w) {
  free(w);
}

Window Window_find(Node head, int key) {
  Node pred = head;
  Node curr = Node_getNext(head);
  while (curr->key < key) {
    pred = curr;
    curr = Node_getNext(pred);
  }
  return Window_create(pred, curr);
}

// =============== End Window class ===============

// =============== Start BucketList class ===============

typedef struct BucketList {
  Node head;
} * BucketList;

BucketList BucketList_create() {
  BucketList this = malloc(sizeof(struct BucketList));
  this->head = Node_create_sentinel(0);
  AtomicMarkableReference_destroy(this->head->next);
  this->head->next =
    AtomicMarkableReference_create(Node_create_sentinel(INT_MAX), false);
  return this;
}

void BucketList_destroy(BucketList this) {
  Node curr = this->head;
  while (curr != NULL) {
    Node next = AtomicMarkableReference_getReference(curr->next);
    Node_destroy(curr);
    curr = next;
  }
  free(this);
}

// second constructor, maybe improve how this is denoted?
BucketList BucketList_BucketList(BucketList this, Node e) {
  this->head = e;
  return this;
}

// forward declarations for bit manipulation functions
int BucketList_reverse(int key);
int BucketList_makeRegularKey(T x);
int BucketList_makeSentinelKey(int key);

bool BucketList_add(BucketList this, T x) {
  int key = BucketList_makeRegularKey(x);
  bool splice;
  while (true) {
    // find predecessor and current entries
    Window window = Window_find(this->head, key);
    Node pred = window->pred;
    Node curr = window->curr;
    // is the key present?
    if (curr->key == key) {
      return false;
    } else {
      // splice in new entry
      Node entry = Node_create(key, x);
      AtomicMarkableReference_set(entry->next, curr, false);
      splice = AtomicMarkableReference_compareAndSet
        (pred->next, curr, entry, false, false);
      if (splice)
        return true;
      else
        continue;
    }
  }
}

bool BucketList_remove(BucketList this, T x) {
  int key = BucketList_makeRegularKey(x);
  bool snip;
  while (true) {
    // find predecessor and current entries
    Window window = Window_find(this->head, key);
    Node pred = window->pred;
    Node curr = window->curr;
    // is the key present?
    if (curr->key != key) {
      return false;
    } else {
      // snip out matching entry
      snip = AtomicMarkableReference_attemptMark(pred->next, curr, true);
      if (snip)
        return true;
      else
        continue;
    }
  }
}

bool BucketList_contains(BucketList this, T x) {
  int key = BucketList_makeRegularKey(x);
  Window window = Window_find(this->head, key);
  Node pred = window->pred;
  Node curr = window->curr;
  return (curr->key == key);
}

BucketList BucketList_getSentinel(BucketList this, int index) {
  int key = BucketList_makeSentinelKey(index);
  bool splice;
  while (true) {
    // find predecessor and current entries
    Window window = Window_find(this->head, key);
    Node pred = window->pred;
    Node curr = window->curr;
    // is the key present?
    if (curr->key == key) {
      return BucketList_BucketList(BucketList_create(), curr);
    } else {
      // splice in new entry
      Node entry = Node_create_sentinel(key);
      AtomicMarkableReference_set
        (entry->next, AtomicMarkableReference_getReference(pred->next), false);
      splice = AtomicMarkableReference_compareAndSet
        (pred->next, curr, entry, false, false);
      if (splice)
        return BucketList_BucketList(BucketList_create(), entry);
      else
        continue;
    }
  }
}

int BucketList_reverse(int key) {
  int loMask = LO_MASK;
  int hiMask = HI_MASK;
  int result = 0;
  for (int i = 0; i < WORD_SIZE; i++) {
    if ((key & loMask) != 0) { // bit set
      result |= hiMask;
    }
    loMask <<= 1;
    hiMask >>= 1;  // fill with 0 from left
  }
  return result;
}

int BucketList_makeRegularKey(T x) {
  int code = hash_code(x) & MASK; // take 3 lowest bytes
  return BucketList_reverse(code | HI_MASK);
}

int BucketList_makeSentinelKey(int key) {
  return BucketList_reverse(key & MASK);
}

// =============== End BucketList class ===============

// =============== Start LockFreeHashSet class ===============

struct Set {
  BucketList* bucket;
  AtomicInteger bucketSize;
  AtomicInteger setSize;
};

Set Set_create() {
  Set this = malloc(sizeof(struct Set));
  this->bucket = malloc(sizeof(BucketList) * CAPACITY);
  this->bucket[0] = BucketList_create();
  for (int i = 1; i < CAPACITY; i++) this->bucket[i] = NULL;
  this->bucketSize = AtomicInteger_create(2);
  this->setSize = AtomicInteger_create(0);
  return this;
}

void Set_destroy(Set this) {}

void Set_initialize_context(void) {}

void Set_finalize_context(void) {}

void Set_initialize(Set set) {}

void Set_finalize(Set set) {}

void Set_terminate(int tid) {}

bool Set_stuck(void) {
  return false;
}

BucketList Set_getBucketList(Set this, int myBucket);

void Set_initializeBucket(Set this, int myBucket);

int Set_getParent(Set this, int myBucket);

bool Set_add(Set this, T x) {
  int myBucket = hash_code(x) % AtomicInteger_get(this->bucketSize);
  BucketList b = Set_getBucketList(this, myBucket);
  if (!BucketList_add(b, x))
    return false;
  int setSizeNow = AtomicInteger_getAndIncrement(this->setSize);
  int bucketSizeNow = AtomicInteger_get(this->bucketSize);
  if (setSizeNow / (double)bucketSizeNow > THRESHOLD)
    AtomicInteger_compareAndSet(this->bucketSize, bucketSizeNow,
                                2 * bucketSizeNow);
  return true;
}

bool Set_remove(Set this, T x) {
  int myBucket = abs(hash_code(x) % AtomicInteger_get(this->bucketSize));
  BucketList b = Set_getBucketList(this, myBucket);
  if (!BucketList_remove(b, x)) {
    return false;
  }
  return true;
}

bool Set_contains(Set this, T x) {
  int myBucket = abs(hash_code(x) % AtomicInteger_get(this->bucketSize));
  BucketList b = Set_getBucketList(this, myBucket);
  return BucketList_contains(b, x);
}

BucketList Set_getBucketList(Set this, int myBucket) {
  if (this->bucket[myBucket] == NULL)
    Set_initializeBucket(this, myBucket);
  return this->bucket[myBucket];
}

void Set_initializeBucket(Set this, int myBucket) {
  int parent = Set_getParent(this, myBucket);
  if (this->bucket[parent] == NULL)
    Set_initializeBucket(this, parent);
  BucketList b = BucketList_getSentinel(this->bucket[parent], myBucket);
  if (b != NULL)
    this->bucket[myBucket] = b;
}

int Set_getParent(Set this, int myBucket) {
  int parent = AtomicInteger_get(this->bucketSize);
  do {
    parent = parent >> 1;
  } while (parent > myBucket);
  parent = myBucket - parent;
  return parent;
}


// Rewrite this. Two problems:
// - Calling Set_getBucketList() on a NULL BucketList creates a new BucketList,
//     just read this->bucket directly instead and skip any NULLs found.
// - Create and use BucketList_print() and Node_print() functions for readability.
void Set_print(Set this) {
  for (int i = 0; i < AtomicInteger_get(this->bucketSize) - 1; i++) {
    BucketList b = Set_getBucketList(this, i);
    Node curr = b->head;
    $print("[", i, "] -> ");
    curr = Node_getNext(curr);
    while (curr != Set_getBucketList(this, i + 1)->head && curr != NULL) {
      if (curr->key != INT_MAX) {
        $print("(", curr->value, ") -> ");
      } else {
        $print("(", curr->key, ") -> ");
      }
      curr = Node_getNext(curr);
    }
  }
  BucketList b = Set_getBucketList(this,
                                   AtomicInteger_get(this->bucketSize) - 1);
  Node curr = b->head;
  $print("[", AtomicInteger_get(this->bucketSize) - 1, "] -> ");
  curr = Node_getNext(curr);
  while (curr != NULL) {
    if (curr->key != INT_MAX) {
      $print("(", curr->value, ") -> ");
    } else {
      $print("(", curr->key, ")");
    }
    curr = Node_getNext(curr);
  }
}


// =============== End LockFreeHashSet class ===============

#ifdef _LOCK_FREE_HASH_SET_MAIN
int main() {
  Set s = Set_create();
  int nthread = 3;
  tid_init(nthread);
  Set_initialize_context();
  Set_initialize(s);
  $parfor (int i : 0 .. nthread-1) {
    tid_register(i);
    Set_add(s, i);
    Set_terminate(i);
    tid_unregister();
  }
  Set_finalize(s);
  Set_finalize_context();
  tid_finalize();

  $print("Set after adding 0..", nthread-1, ": ");
  Set_print(s);
  $print("\n");

  tid_init(1);
  Set_initialize_context();
  Set_initialize(s);
  tid_register(0);
  $for (int i : 0 .. nthread-1)
    $assert(Set_contains(s, i));
  Set_terminate(0);
  tid_unregister();
  Set_finalize(s);
  Set_finalize_context();
  tid_finalize();

  tid_init(nthread);
  Set_initialize_context();
  Set_initialize(s);
  $parfor (int i : 1 .. 2) {
    tid_register(i);
    Set_remove(s, i);
    Set_terminate(i);
    tid_unregister();
  }
  Set_finalize(s);
  Set_finalize_context();
  tid_finalize();

  $print("Set after removing 1..2: ");
  Set_print(s);
  $print("\n");

  $for (int i : 0 .. nthread-1) {
    if (i == 1 || i == 2) {
      $assert(!Set_contains(s, i));
    } else {
      $assert(Set_contains(s, i));
    }
  }

  tid_init(nthread);
  Set_initialize_context();
  Set_initialize(s);
  $parfor (int i : 0 .. nthread-1) {
    tid_register(i);
    if (i == 2)
      Set_add(s, i);
    else
      Set_remove(s, i);
    Set_terminate(i);
    tid_unregister();
  }
  Set_finalize(s);
  Set_finalize_context();
  tid_finalize();

  $print("Set after adding 2 and removing everything else: ");
  Set_print(s);
  $print("\n");

  $for (int i : 0 .. nthread-1) {
    if (i == 2) {
      $assert(Set_contains(s, i));
    } else {
      $assert(!Set_contains(s, i));
    }
  }
}
#endif