mte_helper.c

/*
 * ARM v8.5-MemTag Operations
 *
 * Copyright (c) 2020 Linaro, Ltd.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, see <http://www.gnu.org/licenses/>.
 */

#include "qemu/osdep.h"
#include "cpu.h"
#include "internals.h"
#include "exec/exec-all.h"
#include "exec/ram_addr.h"
#include "exec/cpu_ldst.h"
#include "exec/helper-proto.h"
#include "qapi/error.h"
#include "qemu/guest-random.h"
#include <math.h>

// These are in the C++ standard lib, not available in C:
// #include <map>
// #include <vector>

// std::map<int, std::vector<char> > per_page_tags_v1;
// std::map<page_id, tree_rle_DS> per_page_tags_v2;

// char* btree_global_mem; / /if this is null, alloc a big block of fixed_sz.

// #include <sys/queue.h>
// #include <sys/tree.h>

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <math.h>

// v2 BTree impl begins:
// #define BT_MAX_CHILD 5
// #define BT_MAX_NODE 4
// #define BT_MIN_CHILD 3
// #define BT_MIN_NODE 2

#define REUSING_FREED_NODE 1
#define ALLOC_FREED_ARR_SZ 2

char* global_btree_buffer; // 64-bit METADATA
long long int global_btree_buffer_first_avail; // 64-bit METADATA
// long long int all_time_max_buf_first_avail;
int global_btree_root_offset; //  METADATA
int alloc_freed_nodes; //  METADATA: this is b/w [0,ALLOC_FREED_ARR_SZ]
int alloc_freed_nodes_arr[ALLOC_FREED_ARR_SZ]; // METADATA
int total_alloc_freed_nodes; // [Not metadata] this counts the total #freed nodes

long long int global_stg_call_count;
bool btree_verify_failed = false;
long long int last_printed_global_btree_buf_first_avail, last_printed_total_alloc_freed_nodes;
// long long int min_xlat = 0;
// long long int max_xlat = 0;

#define PROFILE_TIMING_ARR_SZ 100

double stg_common_times_sum = 0.0, mte_probe1_common_times_sum = 0;
long long int stg_count = 0, mte_probe1_count = 0;
double stg_common_times_arr[PROFILE_TIMING_ARR_SZ], stg_tag_arr_times_arr[PROFILE_TIMING_ARR_SZ], stg_btree_times_arr[PROFILE_TIMING_ARR_SZ];
double mte_probe1_common_times_arr[PROFILE_TIMING_ARR_SZ], mte_probe1_tag_arr_times_arr[PROFILE_TIMING_ARR_SZ], mte_probe1_btree_times_arr[PROFILE_TIMING_ARR_SZ];

// total 43 bytes, rounded off to 44.
// struct BTreeNode {
//   int keys[BT_MAX_NODE]; // stores upto 4 32-bit keys : 16byt
//   unsigned char tags[BT_MAX_NODE/2]; // stores a 4bit tag per key : 2byt
//   char count; // 8bits: bits[3:0]: #keys : 1byt
//   int children[BT_MAX_CHILD]; // stores upto 5 32-bit offsets. : 20byt
//   int parent; // stores 32-bit parent ptr offset. : 4byt
// };

void printIntArr(int* arr, int len)
{
  printf("IntArr L:%i:", len);
  for (int x = 0; x < len; x++)
    printf("%i, ", arr[x]);
  // printf("\n");
}

void printCharArr(unsigned char* arr, int len)
{
  for (int x = 0; x < len; x++)
    printf("%i, ", arr[x]); // a[%i] : 
  // printf("\n");
}

void printBTreeNode(struct BTreeNode* node)
{
  printf("N off: %li , p_off: %i , ct: %i | K: ", ((char*)node-global_btree_buffer), node->parent, node->count);
  printIntArr(node->keys, BT_MAX_NODE);
  printf(" | T: ");
  printCharArr(node->tags, BT_MAX_NODE/2);
  printf(" | Ch: ");
  printIntArr(node->children, BT_MAX_CHILD);
}

void printBTree(struct BTreeNode* rnode, int level)
{
  printf("LEVEL %i : ", level);
  if (level < 7)
  {
    printBTreeNode(rnode);
    for (int i = 0; i <= rnode->count; i++)
      if (rnode->children[i] != -1)
        printBTree((struct BTreeNode*)(global_btree_buffer + rnode->children[i]), level+1);    
  }
  else
    printf("ERRORRRRR!!!! TREE DEPTH MORE THAN 6!!!! \n");
}

static struct BTreeNode* getGlobalRootPtr(void)
{
  return (struct BTreeNode*) (global_btree_buffer + global_btree_root_offset);
}

void initBTreeNode(struct BTreeNode* rnode, int parent_offset)
{
  rnode->count = 0;
  for (int i = 0; i < BT_MAX_NODE; i++)
    rnode->keys[i] = -1;
  for (int i = 0; i < BT_MAX_CHILD; i++)
    rnode->children[i] = -1; // null ptr
  rnode->parent = parent_offset;
}

void initBTree(char* global_btree_buffer)
{
  // make a new BTreeNode, with key=0,tag=0. []
  struct BTreeNode* root = (struct BTreeNode*) global_btree_buffer;
  initBTreeNode(root, -1); // parent = null ptr
  root->count = 1;
  root->keys[0] = 0;

  global_btree_buffer_first_avail += sizeof(struct BTreeNode);
  global_btree_root_offset = 0;
  alloc_freed_nodes = 0; // tracks fragmentation.
  total_alloc_freed_nodes = 0;
  for (int i = 0; i < ALLOC_FREED_ARR_SZ; i++)
    alloc_freed_nodes_arr[i] = -1; // initialize freed addresses with -1.

  printf("##### FINISHED initBTree: root ptr: %p, new first_avail: %lli, global_btree_root_offset: %i, alloc_freed_nodes: %i, total_alloc_freed_nodes: %i \n", root, global_btree_buffer_first_avail, global_btree_root_offset, alloc_freed_nodes, total_alloc_freed_nodes);
}

unsigned char getArrIthTag(unsigned char* tag_arr, int i)
{
  return (i%2 == 1) ? ((tag_arr[i/2])%16) : (tag_arr[i/2] >> 4) ;
}

void setArrIthTag(unsigned char* tag_arr, int i, unsigned char tag)
{
  if (tag >= 16)
  {
    printf("ERRORRRRR! setArrIthTag called with invalid tag value!!! %i \n", tag);
    // throw;
  }

  if (i%2 == 1)
    tag_arr[i/2] = (tag_arr[i/2] & 0xf0) | tag;
  else
    tag_arr[i/2] = (tag_arr[i/2] & 0x0f) | (tag << 4);
  // printf("Done with setArrIthTag [inputs: i: %i, tag: %i] new i/2th tag byte: %i \n", i, tag, tag_arr[i/2]);
}

// leaf if all children are -1 (i.e. null ptr)
bool isBTreeNodeLeaf(struct BTreeNode* root)
{
  for (int i = 0; i < BT_MAX_CHILD; i++)
    if (root->children[i] != -1)
      return false;
  return true;
}

// get ptr (in terms of offset from root addr) to node containing largest key <= addr.
// note that since this func is called for random ptrs in BTree, the root input is NOT the absolute root of tree.
// Hence, child ptr = global buffer ptr + offset.
int getLargestNodeLessThan(struct BTreeNode* root, int addr)
{
  if (!root)
    return -1;

  int i=0;
  while ((i < root->count) && (root->keys[i] < addr) )
    i++;

  // case1: i = root->count: ans is in the right most child OR root node.
  if (i >= root->count)
  {
    if ( (i < BT_MAX_CHILD) && (root->children[i] != -1) )
    {
      int ans_in_right_child = getLargestNodeLessThan( (struct BTreeNode*)(global_btree_buffer + root->children[i]), addr );
      return (ans_in_right_child == -1) ? ((char*)root - global_btree_buffer) : ans_in_right_child;
    }
    else
      return ((char*)root - global_btree_buffer); // the current root node has the largest key<=addr.
  }

  // case2: keys[i] = addr: return root node!
  if (root->keys[i] == addr)
    return ((char*)root - global_btree_buffer);

  // case3: keys[i] > addr: go to the just-left child of ith key, i.e. children[i] OR root node.
  if (root->keys[i] > addr)
  {
    int ans_in_left_child = -1;
    if ( (i < BT_MAX_CHILD) && (root->children[i] != -1) )
      ans_in_left_child = getLargestNodeLessThan( (struct BTreeNode*)(global_btree_buffer + root->children[i]), addr );

    return (ans_in_left_child == -1) ? ( (i == 0) ? -1 : ((char*)root - global_btree_buffer) ) : ans_in_left_child;
  }
  return -1;
}

// get ptr (in terms of offset from root addr) to node containing smallest key >= addr.
// note that since this func is called for random ptrs in BTree, the root input is NOT the absolute root of tree.
// Hence, child ptr = global buffer ptr + offset.
int getSmallestNodeMoreThan(struct BTreeNode* root, int addr)
{
  if (!root)
    return -1;

  int i=(root->count)-1;
  while ( (i >= 0) && (root->keys[i] >= addr))
    i--;

  // i < (root->ct)-1 implies there's atleast one elem in root which is >= addr.
  // case1: i = -1, i.e. all elems in array are >= addr
  if (i <= -1)
  {
    if (root->children[0] != -1)
    {
      int ans_in_left_child = getSmallestNodeMoreThan( (struct BTreeNode*)(global_btree_buffer+root->children[0]) , addr);
      return (ans_in_left_child == -1) ? ((char*) root - global_btree_buffer) : ans_in_left_child;
    }
    else
      return ((char*)root - global_btree_buffer); // the current root node has the smallest key>=addr.
  }
  // case2:
  if (root->keys[i] == addr)
    return ((char*)root - global_btree_buffer);

  // case3: all elems including keys[i] are < addr. check right child of keys[i]
  if (root->keys[i] < addr)
  {
    int ans_in_right_child = -1;
    if (root->children[i+1] != -1)
      ans_in_right_child = getSmallestNodeMoreThan((struct BTreeNode*)(global_btree_buffer+root->children[i+1]) , addr);

    return (ans_in_right_child == -1) ? ((i == (root->count-1)) ? -1 : ((char*) root - global_btree_buffer) ) : ans_in_right_child;
  }
  return -1;
}

void getLargestKeyTagLessThan(struct BTreeNode* root, int addr, int* x_key, int* x_tag)
{
  int i=0;
  while ((i < root->count) && (root->keys[i] < addr) )
    i++;

  if (i >= root->count)
  {
    *x_key = root->keys[i-1];
    *x_tag = getArrIthTag(root->tags, i-1);
  }
  else if (root->keys[i] == addr)
  {
    *x_key = root->keys[i];
    *x_tag = getArrIthTag(root->tags, i);
  }
  else if (root->keys[i] > addr)
  {
    if (i == 0)
    {
      // printf("IN getLargestKeyTagLessThan: 1st Elem in root is > addr!!! \n");
      *x_key = -1;
      *x_tag = -1;
      // throw;
    }
    *x_key = root->keys[i-1];
    *x_tag = getArrIthTag(root->tags, i-1);
  }
  else
  {
    printf("IN getLargestKeyTagLessThan: WEIRD ERROR!!!! \n");
    // throw;
  }  
}

void getSmallestKeyTagMoreThan(struct BTreeNode* root, int addr, int* x_key, int* x_tag)
{
  int i=(root->count)-1;
  while ((i >= 0) && (root->keys[i] > addr) )
    i--;

  if (i < 0) // all elems in root are >= addr
  {
    *x_key = root->keys[0];
    *x_tag = getArrIthTag(root->tags, 0);
  }
  else if (root->keys[i] == addr)
  {
    *x_key = root->keys[i];
    *x_tag = getArrIthTag(root->tags, i);
  }
  else if (root->keys[i] < addr)
  {
    if (i == (root->count - 1))
    {
      // printf("IN getSmallestKeyTagMoreThan: last elem is < addr!!! \n");
      *x_key = -1;
      *x_tag = -1;
      // throw;
    }
    *x_key = root->keys[i+1];
    *x_tag = getArrIthTag(root->tags, i+1);
  }
  else
  {
    printf("IN getSmallestKeyTagMoreThan: WEIRD ERROR!!!! \n");
    // throw;
  }  
}

// adding 1 key,tag & 1 child to node
void insertBTreeKeyInNode(struct BTreeNode* rnode, int key, int tag, bool right_child, int child_offset)
{
  // find where to insert
  int i = 0;
  while ((i < rnode->count) && (rnode->keys[i] < key) )
    i++;

  // insert key at rnode->keys[i]
  // shift tags,keys,children
  int curr = key;
  int curr_tag = tag;
  int curr_child = child_offset;
  while (i <= rnode->count)
  {
    int new_curr = rnode->keys[i];
    int new_curr_tag = getArrIthTag(rnode->tags, i);
    int new_curr_child = (right_child) ? rnode->children[i+1] : rnode->children[i];

    rnode->keys[i] = curr;
    setArrIthTag(rnode->tags, i, curr_tag);
    if (right_child)
      rnode->children[i+1] = curr_child;
    else
      rnode->children[i] = curr_child;

    curr = new_curr;
    curr_tag = new_curr_tag;
    curr_child = new_curr_child;

    i++;
  }
  if (!right_child)
    rnode->children[i] = curr_child;

  if (child_offset != -1)
  {
    struct BTreeNode* child_ptr = (struct BTreeNode*) (global_btree_buffer + child_offset);
    child_ptr->parent = (char*)rnode - global_btree_buffer;
  }

  rnode->count += 1;
}

// deletes key&tag from rnode - i.e. updates keys, count, children.
// equivalent to removeVal func in programiz code.
// rebalanceBTree calls this for non-leaf nodes, so,
// del_child: if 0, delete left child of key, else delete right child of key.
bool deleteBTreeKeyFromNode(struct BTreeNode* rnode, int key, int del_child)
{
  int i = 0;
  while ((i < rnode->count) && (rnode->keys[i] != key))
    i++;

  while (i < (rnode->count-1))
  {
    // update tags, keys, children:
    rnode->keys[i] = rnode->keys[i+1];
    setArrIthTag(rnode->tags, i, getArrIthTag(rnode->tags, i+1));
    // children:
    if (del_child == 0)
      rnode->children[i] = rnode->children[i+1];
    else
      rnode->children[i+1] = rnode->children[i+2];
    i++;
  }
  rnode->keys[rnode->count - 1] = -1; // just to avoid confusion when printing/reading BTree node
  // set last tag to 0
  setArrIthTag(rnode->tags, i, 0);
  // move last child, extra move if del_left_child.
  if (del_child == 0)
    rnode->children[i] = rnode->children[i+1];

  rnode->count -= 1;

  return true;
}


// v2: get tag of addr granule, equiv to TagArr[addr]
// NOTE: ALWAYS CALL this func with root as the global root ptr.
unsigned char getBTreeTag(struct BTreeNode* root, int addr)
{
  // struct BTreeNode* groot = (struct BTreeNode*) (global_btree_buffer + global_btree_root_offset);
  if (!root)
    return -1;

  int node_with_largest_less_than_off = getLargestNodeLessThan(getGlobalRootPtr(), addr);
  struct BTreeNode* node_with_largest_less_than = (struct BTreeNode*) (global_btree_buffer + node_with_largest_less_than_off);

  // find the largest key <= addr in the tree
  int i=0;
  while ((i < node_with_largest_less_than->count) && (node_with_largest_less_than->keys[i] < addr) )
    i++;

  // case1: i = root->count: go to the right most child OR Null.
  if (i >= node_with_largest_less_than->count)
  {
    unsigned char node_last_tag = getArrIthTag(node_with_largest_less_than->tags, i-1);
    return node_last_tag;
  }

  // case2: keys[i] = addr: return ith tag!
  if (node_with_largest_less_than->keys[i] == addr)
    return getArrIthTag(node_with_largest_less_than->tags, i); // (i%2 == 1) ? ((root->tags[i/2])%16) : (root->tags[i/2] >> 4); // &(root->tags[i]);

  // case3: keys[i] > addr: go to the just-left child of ith key, i.e. children[i]
  if (node_with_largest_less_than->keys[i] > addr)
  {
    // TODO: its possible that all elems in this child are > addr.
    if (i == 0)
    {
      printf("ERRORRR: No ELEMENT in tree <= addr!!!");
      // throw;
    }
    unsigned char node_last_tag = getArrIthTag(node_with_largest_less_than->tags, i-1);
    return node_last_tag;
    // Old code: 
    // char leftC_tag = getTag(root->children[i], addr);
    // return (leftC_tag != -1) ? leftC_tag : node_last_tag;
  }

  return -1;
}

// re-balance, starting from a deficient leaf node:
void rebalanceBTree(struct BTreeNode* rnode)
{
  struct BTreeNode* rleaf = rnode;
  int iter_ct = 0;
  while (iter_ct < 20)
  {
    int rleaf_offset = (char*) rleaf - global_btree_buffer;
    
    // printf("IN rebalanceBTree: for node %p, parent: %i \n", rleaf, rleaf->parent);

    if (rleaf->parent != -1)
    {
      struct BTreeNode* rleaf_parent = (struct BTreeNode*) (global_btree_buffer + rleaf->parent);
      int i = 0;
      while ((i <= rleaf_parent->count) && (rleaf_parent->children[i] != rleaf_offset))
        i++;

      struct BTreeNode* rleaf_left_bro = NULL;
      struct BTreeNode* rleaf_right_bro = NULL;

      int rleaf_left_bro_offset = -1; // , rleaf_right_bro_offset = -1;

      // check right sibling:
      if ((i < rleaf_parent->count) && (rleaf_parent->children[i+1] != -1))
      {
        // rleaf_right_bro_offset = rleaf_parent->children[i+1];
        rleaf_right_bro = (struct BTreeNode*) (global_btree_buffer + rleaf_parent->children[i+1]);
        if (rleaf_right_bro->count > BT_MIN_NODE)
        {
          // separator key moves to rleaf
          rleaf->keys[(int)(rleaf->count)] = rleaf_parent->keys[i];
          setArrIthTag(rleaf->tags, rleaf->count, getArrIthTag(rleaf_parent->tags, i));
          rleaf->count += 1;

          // use rleaf_right_bro->children[0] as the last child of rleaf:
          rleaf->children[(int)(rleaf->count)] = rleaf_right_bro->children[0];
          if (rleaf_right_bro->children[0] != -1)
          {
            struct BTreeNode* rl_right_bro_c0 = (struct BTreeNode*) (global_btree_buffer + rleaf_right_bro->children[0]);
            rl_right_bro_c0->parent = rleaf_offset;
          }

          // delete left most key from rleaf_right_bro, 
          // use it to replace separator_key in rleaf_parent.
          rleaf_parent->keys[i] = rleaf_right_bro->keys[0];
          setArrIthTag(rleaf_parent->tags, i, getArrIthTag(rleaf_right_bro->tags, 0));
          deleteBTreeKeyFromNode(rleaf_right_bro, rleaf_right_bro->keys[0], 0); // ensures we delete left child of key[0], i.e. child[0]
         
          // printf("IN rebalanceBTree: rotated with rleaf_right_bro: %p (offset: %i) \n", rleaf_right_bro, rleaf_right_bro_offset);
          // printBTreeNode(rleaf);
          // printBTreeNode(rleaf_right_bro);
          // printBTreeNode(rleaf_parent);

          return;
        }
      }
      // check if can rotate with left sibling:
      if ((i > 0) && (rleaf_parent->children[i-1] != -1))
      {
        rleaf_left_bro_offset = rleaf_parent->children[i-1];
        rleaf_left_bro = (struct BTreeNode*) (global_btree_buffer + rleaf_parent->children[i-1]);
        if (rleaf_left_bro->count > BT_MIN_NODE)
        {
          // separator key moves to rleaf TODO: this will move at rleaf->keys[0], slide all existing keys/children/tags:
          insertBTreeKeyInNode(rleaf, rleaf_parent->keys[i-1], getArrIthTag(rleaf_parent->tags,i-1), 0, rleaf_left_bro->children[(int)(rleaf_left_bro->count)]);
          // rleaf->keys[rleaf->count] = rleaf_parent->keys[i-1];
          // setArrIthTag(rleaf->tags, rleaf->count, getArrIthTag(rleaf_parent->tags, i-1));
          // rleaf->count += 1;

          // make rleaf_left_bro's last child, the first child of rleaf: [Doing this as part of insertBTreeKeyInNode]
          // rleaf->children[0] = rleaf_left_bro->children[rleaf_left_bro->count];

          // delete rightmost key from rleaf_left_bro, use it to replace separator_key in rleaf_parent.
          rleaf_parent->keys[i-1] = rleaf_left_bro->keys[rleaf_left_bro->count-1];
          setArrIthTag(rleaf_parent->tags, i-1, getArrIthTag(rleaf_left_bro->tags, rleaf_left_bro->count-1));
          deleteBTreeKeyFromNode(rleaf_left_bro, rleaf_left_bro->keys[rleaf_left_bro->count-1], 1 ); // deleting right child of last key

          // printf("IN rebalanceBTree: rotated with rleaf_right_bro: %p", rleaf_left_bro);
          // printBTreeNode(rleaf);
          // printBTreeNode(rleaf_left_bro);
          // printBTreeNode(rleaf_parent);

          return;
        }
      }
      // Need to merge: [no sibling with >min_keys]
      // put separator, rleaf->keys in either left_bro or right_bro.
      if (rleaf_left_bro != NULL)
      {
        // put separator in rleaf_left_bro
        rleaf_left_bro->keys[(int)(rleaf_left_bro->count)] = rleaf_parent->keys[i-1];
        setArrIthTag(rleaf_left_bro->tags, rleaf_left_bro->count, getArrIthTag(rleaf_parent->tags,i-1));
        rleaf_left_bro->count += 1;

        // put rleaf->keys in rleaf_left_bro
        for (int rx = 0; rx < rleaf->count; rx++)
        {
          rleaf_left_bro->keys[(int)(rleaf_left_bro->count)] = rleaf->keys[rx];
          setArrIthTag(rleaf_left_bro->tags, rleaf_left_bro->count, getArrIthTag(rleaf->tags,rx));
          rleaf_left_bro->children[(int)(rleaf_left_bro->count)] = rleaf->children[rx];
          if (rleaf->children[rx] != -1)
          {
            struct BTreeNode* rleaf_child_rx = (struct BTreeNode*) (global_btree_buffer + rleaf->children[rx]);
            rleaf_child_rx->parent = rleaf_left_bro_offset;
          }
          rleaf_left_bro->count += 1;
        }
        rleaf_left_bro->children[(int)(rleaf_left_bro->count)] = rleaf->children[(int)(rleaf->count)];
        if (rleaf->children[(int)(rleaf->count)] != -1)
        {
          struct BTreeNode* rleaf_child_last = (struct BTreeNode*) (global_btree_buffer + rleaf->children[(int)(rleaf->count)]);
          rleaf_child_last->parent = rleaf_left_bro_offset;
        }

        // delete separator from parent: delete right child of this key
        deleteBTreeKeyFromNode(rleaf_parent, rleaf_parent->keys[i-1], 1); // rleaf_parent is a non leaf node!!!

        // printf("IN rebalanceBTree: Merged with rleaf_left_bro: %p \n", rleaf_left_bro);
        // printBTreeNode(rleaf_left_bro);
        // printBTreeNode(rleaf_parent);

      }
      else if (rleaf_right_bro != NULL)
      {
        // Done: Handle children here!! [insertBTreeKeyInNode handles it]
        // children rleaf_right_bro: key_ct += 1+rleaf->ct, child_ct += 0+rleaf->ct+1. Works!

        // put rleaf's keys in rleaf_right_bro
        for (int lx = 0; lx < rleaf->count; lx++)
          insertBTreeKeyInNode(rleaf_right_bro, rleaf->keys[lx], getArrIthTag(rleaf->tags,lx), 0, rleaf->children[lx]);
      
        // put separator in rleaf_right_bro
        insertBTreeKeyInNode(rleaf_right_bro, rleaf_parent->keys[i], getArrIthTag(rleaf_parent->tags, i), 0, rleaf->children[(int)(rleaf->count)]);

        // delete separator from parent: delete left child of this key
        deleteBTreeKeyFromNode(rleaf_parent, rleaf_parent->keys[i], 0);

        // printf("IN rebalanceBTree: Merged with rleaf_right_bro: %p \n", rleaf_right_bro);
        // printBTreeNode(rleaf_right_bro);
        // printBTreeNode(rleaf_parent);

      }
      else
      {
        printf("ERRORRR!!! Both left, right bros are NULL!! \n");
        // throw;
      }
      // 'delete' rleaf node.
      freeNodeOffset(rleaf_offset);
      // Moved this code to freeNodeOffset function:
      // alloc_freed_nodes_arr[total_alloc_freed_nodes%ALLOC_FREED_ARR_SZ] = rleaf_offset;
      // total_alloc_freed_nodes += 1;
      // alloc_freed_nodes = std::min(alloc_freed_nodes+1, ALLOC_FREED_ARR_SZ);
      // // printf("IN rebalanceBTree: UPDATED alloc_freed_nodes to %i \n", alloc_freed_nodes);
      // printIntArr(alloc_freed_nodes_arr, ALLOC_FREED_ARR_SZ);
      // // freeNodeOffset code ends

      // Call re-balance on rleaf_parent, if needed
      if (rleaf_parent->count == 0)
      {
        if (rleaf_parent->parent != -1)
        {
          printf("ERRORRR: SOMEHOW NON-ROOT node is EMPTY!!!!");
          // throw;
        }
        else
        {
          // free current root:
          freeNodeOffset(global_btree_root_offset);
          // alloc_freed_nodes_arr[total_alloc_freed_nodes%ALLOC_FREED_ARR_SZ] = global_btree_root_offset;
          // total_alloc_freed_nodes += 1;
          // printIntArr(alloc_freed_nodes_arr, ALLOC_FREED_ARR_SZ);

          // make rleaf_left/right_bro as the new root.
          struct BTreeNode* new_root = (rleaf_left_bro != NULL) ? rleaf_left_bro : rleaf_right_bro;
          new_root->parent = -1;
          global_btree_root_offset = (char*) new_root - global_btree_buffer;

          // printf("IN rebalanceBTree: JUST MODIFIED the global_btree_root_offset to %i !!!! \n", global_btree_root_offset);
          return;
        }
      }
      if (rleaf_parent->count < BT_MIN_NODE)
        rleaf = rleaf_parent;
      else
        return; // WE ARE DONE!!! [Update: Mar24]
    }
    else
    {
      // rleaf is the root Node. Its okay for it to be deficient.
      // printf("ROOT has %i keys, < min %i \n", rleaf->count, BT_MIN_NODE);
      if (rleaf->count == 0)
      {
        printf("ERRORRR!!!! Somehow ROOT is empty, and we called rebalance!! \n");
        // throw; // this should never happen!
      }
      return; // WE'RE DONE!!!
    }
    iter_ct++;
  }
  printf("ERRORRRRRR!!!!!! RECURSION IN deleteBTreeKey MORE THAN 20 TIMES!!!!");
}


// deletes key from rnode, and tries to balance the BTree
bool deleteBTreeKey(struct BTreeNode* rnode, int key)
{
  if (isBTreeNodeLeaf(rnode))
  {
    deleteBTreeKeyFromNode(rnode, key, 1);
    if (rnode->count < BT_MIN_NODE)
    {
      // re-balance
      rebalanceBTree(rnode);
    }
  }
  else
  {
    int i = 0;
    while ((i < rnode->count) && (rnode->keys[i] != key))
      i++;
    // Need to find a replacement for this key: largest key' < key OR smallest key' > key
    // node with smallest elem >= key.
    int smallest_node_more_key_off = getSmallestNodeMoreThan( (struct BTreeNode*)(global_btree_buffer + rnode->children[i+1]), key );
    struct BTreeNode* smallest_node_more_key = (struct BTreeNode*) (global_btree_buffer + smallest_node_more_key_off);
    // replacement should be in a leaf.
    if (!isBTreeNodeLeaf(smallest_node_more_key))
    {
      printf("IN deleteBTreeKey: ERRORRRR!!!! smallestNodeMoreThan key %i is NOT a leaf!!! \n", key);
      return false;
    }
    // remove the replacement (key,tag) from its leaf and put it in rnode (at ith place)
    int sj = smallest_node_more_key->count-1;
    while ( (sj >= 0) && (smallest_node_more_key->keys[sj] > key) )
      sj--;

    // printf("IN deleteBTreeKey: Deleting %i from node %p, borrowing key %i from leaf_ptr %p \n", key, rnode, smallest_node_more_key->keys[sj+1], smallest_node_more_key);

    rnode->keys[i] = smallest_node_more_key->keys[sj+1]; // this was previously key, replacing with smallest > key.
    setArrIthTag(rnode->tags, i, getArrIthTag(smallest_node_more_key->tags, sj+1) );
    deleteBTreeKeyFromNode(smallest_node_more_key, smallest_node_more_key->keys[sj+1], 1); // note smallest_node_more_key is a leaf

    // printBTreeNode(rnode);
    // printBTreeNode(smallest_node_more_key);

    // if leaf count < BT_MIN_NODE
    if (smallest_node_more_key->count < BT_MIN_NODE)
    {
      // re-balance
      rebalanceBTree(smallest_node_more_key);
    }
  }
  return true;
}

int findLeafNodeToInsertAddr(struct BTreeNode* rnode, int addr)
{
  if (isBTreeNodeLeaf(rnode))
    return (char*)rnode - global_btree_buffer;

  int i = 0;
  while (i < rnode->count && (rnode->keys[i] < addr) )
    i++;

  if ( rnode->children[i] == -1)
  {
    printf("In findLeafNodeToInsertAddr: ERROR!!! Non leaf node (count=%i) has null child (%ith child: %i)!!! \n", rnode->count, i, rnode->children[i]);
    return -1;
    // throw;
  }
  else
    return findLeafNodeToInsertAddr( (struct BTreeNode*)(global_btree_buffer + rnode->children[i]), addr );
  // leftmost / rightmost / middle child:
  // if (i == rnode->count)
  //   return findLeafNodeToInsertAddr( global_btree_buffer + rnode->children[i], addr);
  // else if (i == 0)
  //   return findLeafNodeToInsertAddr( global_btree_buffer + rnode->children[0], addr);
}

// if there's any child that has the wrong parent ptr, returns false.
bool verifyBTreePtrs(struct BTreeNode* root)
{
    int root_offset = (char*) root - global_btree_buffer;
    for (int i = 0; i <= root->count; i++)
    {
        if (root->children[i] != -1)
        {
            struct BTreeNode* r_child_i = (struct BTreeNode*) (global_btree_buffer + root->children[i]);
            if (r_child_i->parent != root_offset)
                return false;
            if (!verifyBTreePtrs(r_child_i))
                return false;
        }
    }
    return true;
}

void freeNodeOffset(int offset)
{
  alloc_freed_nodes_arr[alloc_freed_nodes%ALLOC_FREED_ARR_SZ] = offset;

  total_alloc_freed_nodes += 1;
  alloc_freed_nodes = (alloc_freed_nodes+1) > (ALLOC_FREED_ARR_SZ) ? ALLOC_FREED_ARR_SZ : (alloc_freed_nodes+1); //std::min(alloc_freed_nodes+1, ALLOC_FREED_ARR_SZ);
  // printf("FREE-ed addr %i, alloc_freed_nodes: %i, total_alloc_freed_nodes: %i ", offset, alloc_freed_nodes, total_alloc_freed_nodes);
  // printIntArr(alloc_freed_nodes_arr, ALLOC_FREED_ARR_SZ);
}

// Use this func to get offset for assigning new node:
// Manages free-d nodes' space. [if REUSING_FREED_NODE == 1, it tries to re-use freed node space]
long long int getNewNodeOffset(void)
{
  long long int ans = 0;
  if ((alloc_freed_nodes == 0) || (REUSING_FREED_NODE == 0))
  {
    // Use fresh unused space, increment buffer_first_avail:
    ans = global_btree_buffer_first_avail;
    global_btree_buffer_first_avail += sizeof(struct BTreeNode);
    // return ans;
  }
  else
  {
    // use a free-d addr, update free-d count:
    // printf("IN getNewNodeOffset: RE-using a free-d space: alloc_freed_nodes: %i, total_alloc_freed_nodes: %i", alloc_freed_nodes, total_alloc_freed_nodes);
    // printIntArr(alloc_freed_nodes_arr, ALLOC_FREED_ARR_SZ);

    alloc_freed_nodes = (alloc_freed_nodes - 1);
    total_alloc_freed_nodes -= 1;

    ans = alloc_freed_nodes_arr[alloc_freed_nodes % ALLOC_FREED_ARR_SZ];
    alloc_freed_nodes_arr[alloc_freed_nodes % ALLOC_FREED_ARR_SZ] = -1;
    
    if (ans == -1)
      printf("IN getNewNodeOffset: ERRORRRR!!!! freed_addr was -1!!! alloc_freed_nodes: %i, total_alloc_freed_nodes: %i, addr: %lli", alloc_freed_nodes, total_alloc_freed_nodes, ans);
  }
  return ans;
}

// splits rnode->keys + insert_addr array.
// returns offset for new node, and writes median_key to the ptr, and tag for median key to the ptr.
// ASSUMES rnode has BT_MAX_NODE keys!!!
// Note: need to pass new child ptr too, since we need it to correctly update children of new_node_right that we make here:
int splitNode(struct BTreeNode* rnode, int insert_addr, int insert_tag, int insert_child_right_offset, int* pmedian_key, unsigned char* pmedian_tag)
{
  if (rnode->count < BT_MAX_NODE)
  {
    printf("ERROR: SPLITNODE CALLED for NON-FULL node!!!");
    // throw;
  }
  // find median, split.
    int i = 0; 
    while ( (i < rnode->count) && (rnode->keys[i] < insert_addr) )
      i++;

    // ith elem would be addr, elems i - count will be shifted to right.
    int median_id = (rnode->count + 1)/2;
    int median_key = (median_id < i) ? rnode->keys[median_id] : ( (median_id == i) ? insert_addr : (rnode->keys[median_id-1]) ) ;

    // make new arr with rnode->count + 1 elems. Will make updating rnode, new_node_right MUCH easier.
    int full_arr[rnode->count+1];
    unsigned char full_tag_arr[ (int)ceil((rnode->count+1)/2.0) ]; // Note: important to divide by 2.0, o.w. ct/2 was int division.
    int full_children_arr[rnode->count+2];

    // printf("rnode->count %i, %f, ceil((rnode->count+1)/2): %f, int_cast: %i \n", rnode->count, (rnode->count+1)/2.0, ceil((rnode->count+1)/2.0), (int)ceil((rnode->count+1)/2.0));
    for (int xx = 0; xx < (int)ceil((rnode->count+1)/2.0); xx++)
      full_tag_arr[xx] = 0;
      // printf("full_tag_arr[%i] : %i \n", xx, full_tag_arr[xx]);

    for (int j = 0; j < i; j++)
    {
      full_arr[j] = rnode->keys[j];
      setArrIthTag(full_tag_arr, j, getArrIthTag(rnode->tags, j) );
      full_children_arr[j] = rnode->children[j];
    }
    full_arr[i] = insert_addr;
    setArrIthTag(full_tag_arr, i, insert_tag);
    full_children_arr[i] = rnode->children[i]; // note that this child would've been split in a previous recursive call.
    full_children_arr[i+1] = insert_child_right_offset; // the new_node_right made by previous recursive call
    for (int j = i; j < rnode->count; j++)
    {
      full_arr[j+1] = rnode->keys[j];
      setArrIthTag(full_tag_arr, j+1, getArrIthTag(rnode->tags, j) );
      full_children_arr[j+2] = rnode->children[j+1];
    }

    // update rnode keys & count & children:
    // We DO need to update full_chil_arr[0 - median_id]->parent to rnode?
    // median_id is always <= rnode->count, but we dont know where insert_addr, insert_child_right_offset are in the full_arr!
    for (int j = 0; j < BT_MAX_CHILD; j++)
    {
      if (j<median_id)
      {
        rnode->keys[j] = full_arr[j];
        setArrIthTag(rnode->tags, j, getArrIthTag(full_tag_arr, j) );
        rnode->children[j] = full_children_arr[j];        
      }
      else if (j == median_id)
        rnode->children[j] = full_children_arr[j];
      else
        rnode->children[j] = -1; // null ptrs

      // Update children's parent to rnode, for j <= median_id:
      if ((j <= median_id) && (full_children_arr[j] != -1))
      {
        struct BTreeNode* full_child_j = (struct BTreeNode*) (global_btree_buffer + full_children_arr[j]);
        full_child_j->parent = (char*)rnode - global_btree_buffer;
      }
    }
    rnode->count = median_id;

    // update the return values:
    *pmedian_key = full_arr[median_id];
    *pmedian_tag = getArrIthTag(full_tag_arr, median_id);

    if (median_key != full_arr[median_id])
        printf("SPLITNODE: ERRORRR!!! median key %i doens't match full_arr median %i !! \n", median_key, full_arr[median_id]);

    // Split rnode: elems>median_key in new node.
    int new_node_create_offset = getNewNodeOffset();
    struct BTreeNode* new_node_right = (struct BTreeNode*) (global_btree_buffer + new_node_create_offset);
    initBTreeNode(new_node_right, -1);
    int new_node_right_offset = new_node_create_offset;
    // NOTE: we now do this in getNewNodeOffset func: global_btree_buffer_first_avail += sizeof(struct BTreeNode);
    // // new_node_right->count = rnode->count / 2;
    
    // add keys, tags, children to new_node_right.
    for (int j = 0; j < BT_MAX_CHILD; j++)
    {
      int jth_id = j+median_id+1;
      struct BTreeNode* new_node_right_child = NULL;
      if (full_children_arr[jth_id] != -1)
        new_node_right_child = (struct BTreeNode*) (global_btree_buffer + full_children_arr[jth_id]);

      if (jth_id < (BT_MAX_NODE+1))
      {
        new_node_right->keys[j] = full_arr[jth_id];
        setArrIthTag(new_node_right->tags, (j), getArrIthTag(full_tag_arr, jth_id) );
        new_node_right->children[j] = full_children_arr[jth_id];
        if (new_node_right_child)
          new_node_right_child->parent = new_node_right_offset;
        new_node_right->count += 1;
      }
      else if ( jth_id == (BT_MAX_NODE + 1) )
      {
        new_node_right->children[j] = full_children_arr[jth_id]; // last elem of full_children_arr
        if (new_node_right_child)
          new_node_right_child->parent = new_node_right_offset;
      }
      else
        new_node_right->children[j] = -1; // NULL ptr
    }

    // update new_node_right's parent ptr:
    new_node_right->parent = rnode->parent;

    // printf("SPLIT node ptr %p (%li) [orig ct: %i, new ct: %i] [inserting addr: %i, median_key: %i (full_arr[%i]: %i), new_node_right ptr: %p (%li), count: %i] \n", rnode, (char*)rnode - global_btree_buffer, BT_MAX_NODE, rnode->count, insert_addr, median_key, median_id, full_arr[median_id], new_node_right, (char*)new_node_right - global_btree_buffer, new_node_right->count);
    // printBTreeNode(rnode);
    // printBTreeNode(new_node_right);
    return ((char*)new_node_right - global_btree_buffer);
}

bool insertBTreeKey(int addr, unsigned char tag)
{
  // find the leaf where we'll insert addr:
  int leaf_to_insert_offset = findLeafNodeToInsertAddr(getGlobalRootPtr(), addr);
  struct BTreeNode* leaf_to_insert = (struct BTreeNode*) (global_btree_buffer + leaf_to_insert_offset);

  // printf("IN insertBTreeKey: addr: %i, tag: %i, leaf_to_insert_offset: %i \n", addr, tag, leaf_to_insert_offset);

  // we recurse, and update insert_addr, insert_tag, insert_child_right, node_to_insert in each iteration:
  int insert_addr = addr;
  int insert_tag = tag;
  int insert_child_right_offs = -1;
  struct BTreeNode* node_to_insert = leaf_to_insert;
  while (true)
  {
    if (node_to_insert->count < BT_MAX_NODE)
    {
      // just insert in this leaf & done
      int i = 0; 
      while ( (i < node_to_insert->count) && (node_to_insert->keys[i] < insert_addr) )
        i++;
      
      if (i == node_to_insert->count)
      {
        node_to_insert->keys[i] = insert_addr;
        node_to_insert->children[i+1] = insert_child_right_offs;
        setArrIthTag(node_to_insert->tags, i, insert_tag);
      }
      else
      {
        // move all other elems [keys, tags, children] to right.
        int curr = node_to_insert->keys[i];
        int curr_tag = getArrIthTag(node_to_insert->tags, i);
        int curr_child = node_to_insert->children[i+1];

        // put insert_addr at i:
        node_to_insert->keys[i] = insert_addr;
        setArrIthTag(node_to_insert->tags, i, insert_tag);
        node_to_insert->children[i+1] = insert_child_right_offs;
        i++;
        while (i <= node_to_insert->count)
        {
          int new_curr = node_to_insert->keys[i];
          int new_curr_tag = getArrIthTag(node_to_insert->tags, i);
          int new_curr_child = node_to_insert->children[i+1];
          node_to_insert->keys[i] = curr;
          setArrIthTag(node_to_insert->tags, i, curr_tag);
          node_to_insert->children[i+1] = curr_child;

          curr = new_curr;
          curr_tag = new_curr_tag;
          curr_child = new_curr_child;
          i++;
        }
      }

      // update insert_child_right's parent to node_to_insert.
      if (insert_child_right_offs != -1)
      {
        struct BTreeNode* insert_child_right = (struct BTreeNode*) (global_btree_buffer + insert_child_right_offs);
        insert_child_right->parent = (char*)(node_to_insert) - global_btree_buffer;
      }

      node_to_insert->count += 1;
      return true;
    }
    else
    {
      // get median key, push to parent. Parent might split too, recurse till root.    
      int median_key;
      unsigned char median_tag;
      int new_node_right_offset = splitNode(node_to_insert, insert_addr, insert_tag, insert_child_right_offs, &median_key, &median_tag);

      // recurse to parent:
      if (node_to_insert->parent == -1) // node_to_insert was root
      {
        // Construct new root!
        int new_node_create_offset = getNewNodeOffset();
        struct BTreeNode* new_root = (struct BTreeNode*)(global_btree_buffer + new_node_create_offset);
        initBTreeNode(new_root, -1); // parent = -1
        new_root->keys[0] = median_key;
        setArrIthTag(new_root->tags, 0, median_tag);
        new_root->children[0] = (char*)node_to_insert - global_btree_buffer;
        new_root->children[1] = new_node_right_offset;
        new_root->count = 1;

        // update parents of node_to_insert & new_node_right_offset:
        node_to_insert->parent = new_node_create_offset;
        struct BTreeNode* new_node_right = (struct BTreeNode*)(global_btree_buffer+new_node_right_offset);
        new_node_right->parent = new_node_create_offset;

        // update global_btree_root_offset, first_avail
        global_btree_root_offset = new_node_create_offset; // (char*)new_root - global_btree_buffer;
        // global_btree_buffer_first_avail += sizeof(struct BTreeNode); // NO need to do this cuz we do this in getNewNodeOffset().
        // printf("IN insertBTreeKey: JUST MODIFIED the global_btree_root_offset to %i !!!! \n", global_btree_root_offset);

        return true;
      }
      else
      {
        int node_to_insert_offset = node_to_insert->parent;
        // recurse here:
        insert_addr = median_key;
        insert_tag = median_tag;
        insert_child_right_offs = new_node_right_offset;
        node_to_insert = (struct BTreeNode*) (global_btree_buffer + node_to_insert_offset);

        // printf("IN insertBTreeKey: Recursing for splitNode on node %p , inserting addr %i \n", node_to_insert, insert_addr);
      }
      
    }
  }

  
}

// set tag of addr granule.
// TODO: root is at (global_btree_buffer + global_btree_root_offset)
void setBTreeTag(struct BTreeNode* root, int addr, int tag)
{
  // printf("STARTING setTag root: %p , addr: %i, tag: %i \n", root, addr, tag);
  // struct BTreeNode* groot = (struct BTreeNode*) (global_btree_buffer + global_btree_root_offset);
  // first find largest key x <= addr.
  if (!root)
    return;

  int largest_node_less_offset = getLargestNodeLessThan(getGlobalRootPtr(), addr);
  int tag_x = -1, x = -1, x_ind = -1;
  struct BTreeNode* largest_node_less = NULL;
  if (largest_node_less_offset != -1)
  {
    // get tag_x, x from largest_node_less_offset:
    largest_node_less = (struct BTreeNode*) (global_btree_buffer + largest_node_less_offset);
    getLargestKeyTagLessThan(largest_node_less, addr, &x, &tag_x);

    x_ind = 0;
    while ((x_ind < largest_node_less->count) && (largest_node_less->keys[x_ind] != x))
      x_ind++;
  }
  else
  {
    printf("IN setTag: COULDN'T FIND ANY x <= addr (%i) ERRORR !!! \n", addr);
    // throw;
  }

  // if tag[x] = tag, done
  if (tag_x == tag)
  {
    // printf("IN setTag: [x: %i] Returning cuz current tag for addr %i is %i \n !!!", x, addr, tag_x);
    return;
  }

  // Find the run on immediate right (xnext) i.e. smallest>x, smallest>=x+1. Needed for later logic.
  int smallest_node_more_x_offset = getSmallestNodeMoreThan(getGlobalRootPtr(), x+1);
  int tag_xnext = -1, xnext = -1;
  struct BTreeNode* smallest_node_more_x = NULL;

  int tag_xprev = -1, xprev = -1;

  if (smallest_node_more_x_offset != -1)
  {
    // This offset could be -1!!! (i.e. x is the last run)
    smallest_node_more_x = (struct BTreeNode*) (global_btree_buffer + smallest_node_more_x_offset);
    getSmallestKeyTagMoreThan(smallest_node_more_x, x+1, &xnext, &tag_xnext);    
  }

  // printf("Midway thru setTag: x: %i, tag_x: %i, xnext: %i, tag_xnext: %i \n", x, tag_x, xnext, tag_xnext);

  // Possibility of merging with left run.
  if (x == addr)
  {
    // find xprev: largest key < x. i.e. largest <= (x-1)
    int largest_node_less_x_offset = getLargestNodeLessThan( getGlobalRootPtr(), x-1);

    // printf("IN setTag: largest_node_less_offset: %i \n", largest_node_less_x_offset);

    // This offset could be -1!!! (i.e. x is the first run)
    if (largest_node_less_x_offset != -1)
    {
      struct BTreeNode* largest_node_less_x = (struct BTreeNode*) (global_btree_buffer + largest_node_less_x_offset);
      getLargestKeyTagLessThan(largest_node_less_x, x-1, &xprev, &tag_xprev);

      // printf("IN setTag: xprev: %i , tag_xprev: %i \n", xprev, tag_xprev);

      // if tag = tag[xprev] & addr=x : [merge with xprev run]
      if (tag == tag_xprev)
      {
        // if xnext > (x+1) OR there is no xnext: change x to x+1. 
        if ((xnext > (x+1)) || (xnext == -1))
        {
          // change largest_node_less keys arr: replace x by x+1.
          int i = 0;
          while ((i < largest_node_less->count) && (largest_node_less->keys[i] != x))
            i++;
          largest_node_less->keys[i] = x+1;

          // printf("IN setTag: addr: %i, x: %i, xnext: %i, xprev: %i, tag_xprev: %i [largest_node_less,Returning!] \n", addr, x, xnext, xprev, tag_xprev);
          // printBTreeNode(largest_node_less);

          return;
        }
        // else: delete key x: a. shrink node[keys], b. if node empty: ??
        else if (xnext == (x+1))
        {
          deleteBTreeKey(largest_node_less, x);

          // printf("IN setTag: addr: %i, x: %i, xnext: %i, xprev: %i, tag_xprev: %i [largest_node_less,Returning!] \n", addr, x, xnext, xprev, tag_xprev);
          // printBTreeNode(largest_node_less);

          return;
        }
      }
    }
    
  }
  
  // find xnext: smallest key > x.
  // else if tag = tag[xnext] & addr=xnext-1: [merge with xnext run]
  // if xnext == -1, we won't enter this conditional.
  if ((addr == (xnext - 1) ) && (tag == tag_xnext))
  {
    if (xnext > (x+1))
    {
      // Done: Change xnext value to xnext-1, in smallest_node_more_x:
      int i = smallest_node_more_x->count-1;
      while ( (i >= 0) && (smallest_node_more_x->keys[i] > xnext) )
        i--;
      smallest_node_more_x->keys[i] = xnext-1;

      // printf("IN setTag: addr: %i, x: %i, xnext: %i, xprev: %i, tag_xprev: %i, tag_xnext: %i [smallest_node_more_x,Returning!] \n", addr, x, xnext, xprev, tag_xprev, tag_xnext);
      // printBTreeNode(smallest_node_more_x);

      return;
    }
    else
    {
      // delete x and replace xnext in smallest_node_more_x by xnext-1 = x.
      // Option1: delete x, delete (xnext,tag) & insert (x,tag)
      // THIS: Option2: update x' tag to xnext_tag & delete (xnext,tag)
      
      setArrIthTag(largest_node_less->tags, x_ind, tag_xnext);
      deleteBTreeKey(smallest_node_more_x, xnext);

      // printf("IN setTag: addr: %i, x: %i, xnext: %i, xprev: %i, tag_xprev: %i, tag_xnext: %i [smallest_node_more_x,Returning!] \n", addr, x, xnext, xprev, tag_xprev, tag_xnext);
      // printBTreeNode(smallest_node_more_x);

      return;
    }
  }

  // DONE: else [split x into 3 keys]: x,tag[x] | addr,tag | addr+1, i.e. potentially inserting 2 new keys.
  {
    // printf("IN setTag: Couldn't merge newTag with left/right runs. Splitting x into 3 keys \n");
    // x <= addr
    if (x < addr)
    {
      // insert addr,tag as new key
      insertBTreeKey(addr, tag);      
    }
    else
      setArrIthTag(largest_node_less->tags, x_ind, tag);
    // insert a+1,tag_x as new key, BUT if addr+1 = xnext, then addr+1's tag is actually tag_xnext,
    // we dont need to insert (addr+1,tag) separately.
    if ((xnext == -1) || ((addr+1) < xnext))
      insertBTreeKey(addr+1, tag_x);
  
    // printBTreeNode(largest_node_less);
    // printBTreeNode(getGlobalRootPtr());
  }

}


///////////////////////////////////// ORIGINAL mte_helper.c CODE: 


static int choose_nonexcluded_tag(int tag, int offset, uint16_t exclude)
{
    if (exclude == 0xffff) {
        return 0;
    }
    if (offset == 0) {
        while (exclude & (1 << tag)) {
            tag = (tag + 1) & 15;
        }
    } else {
        do {
            do {
                tag = (tag + 1) & 15;
            } while (exclude & (1 << tag));
        } while (--offset > 0);
    }
    return tag;
}

/**
 * allocation_tag_mem:
 * @env: the cpu environment
 * @ptr_mmu_idx: the addressing regime to use for the virtual address
 * @ptr: the virtual address for which to look up tag memory
 * @ptr_access: the access to use for the virtual address
 * @ptr_size: the number of bytes in the normal memory access
 * @tag_access: the access to use for the tag memory
 * @tag_size: the number of bytes in the tag memory access
 * @ra: the return address for exception handling
 *
 * Our tag memory is formatted as a sequence of little-endian nibbles.
 * That is, the byte at (addr >> (LOG2_TAG_GRANULE + 1)) contains two
 * tags, with the tag at [3:0] for the lower addr and the tag at [7:4]
 * for the higher addr.
 *
 * Here, resolve the physical address from the virtual address, and return
 * a pointer to the corresponding tag byte.  Exit with exception if the
 * virtual address is not accessible for @ptr_access.
 *
 * The @ptr_size and @tag_size values may not have an obvious relation
 * due to the alignment of @ptr, and the number of tag checks required.
 *
 * If there is no tag storage corresponding to @ptr, return NULL.
 */
static uint8_t *allocation_tag_mem(CPUARMState *env, int ptr_mmu_idx,
                                   uint64_t ptr, MMUAccessType ptr_access,
                                   int ptr_size, MMUAccessType tag_access,
                                   int tag_size, uintptr_t ra, int* bt_tag_addr)
{
#ifdef CONFIG_USER_ONLY
    uint64_t clean_ptr = useronly_clean_ptr(ptr);
    int flags = page_get_flags(clean_ptr);
    uint8_t *tags;
    uintptr_t index;

    if (!(flags & (ptr_access == MMU_DATA_STORE ? PAGE_WRITE : PAGE_READ))) {
        /* SIGSEGV */
        arm_cpu_tlb_fill(env_cpu(env), ptr, ptr_size, ptr_access,
                         ptr_mmu_idx, false, ra);
        g_assert_not_reached();
    }

    /* Require both MAP_ANON and PROT_MTE for the page. */
    if (!(flags & PAGE_ANON) || !(flags & PAGE_MTE)) {
        return NULL;
    }

    // ADITI comments:
    // page->target data is the tag array. Initialized in page_alloc_target_data. 
    // alloc_size: ~ #granules?
    // so tags init is an arr of sz alloc_size.
    tags = page_get_target_data(clean_ptr);
    if (tags == NULL) {
        size_t alloc_size = TARGET_PAGE_SIZE >> (LOG2_TAG_GRANULE + 1);
        tags = page_alloc_target_data(clean_ptr, alloc_size);
        assert(tags != NULL);
    }

    index = extract32(ptr, LOG2_TAG_GRANULE + 1,
                      TARGET_PAGE_BITS - LOG2_TAG_GRANULE - 1);
    return tags + index;
#else
    uintptr_t index;
    CPUIOTLBEntry *iotlbentry;
    int in_page, flags;
    ram_addr_t ptr_ra;
    hwaddr ptr_paddr, tag_paddr, xlat;
    MemoryRegion *mr;
    ARMASIdx tag_asi;
    AddressSpace *tag_as;
    void *host;

    /*
     * Probe the first byte of the virtual address.  This raises an
     * exception for inaccessible pages, and resolves the virtual address
     * into the softmmu tlb.
     *
     * When RA == 0, this is for mte_probe1.  The page is expected to be
     * valid.  Indicate to probe_access_flags no-fault, then assert that
     * we received a valid page.
     */
    flags = probe_access_flags(env, ptr, ptr_access, ptr_mmu_idx,
                               ra == 0, &host, ra);
    assert(!(flags & TLB_INVALID_MASK));

    /*
     * Find the iotlbentry for ptr.  This *must* be present in the TLB
     * because we just found the mapping.
     * TODO: Perhaps there should be a cputlb helper that returns a
     * matching tlb entry + iotlb entry.
     */
    index = tlb_index(env, ptr_mmu_idx, ptr);
# ifdef CONFIG_DEBUG_TCG
    {
        CPUTLBEntry *entry = tlb_entry(env, ptr_mmu_idx, ptr);
        target_ulong comparator = (ptr_access == MMU_DATA_LOAD
                                   ? entry->addr_read
                                   : tlb_addr_write(entry));
        g_assert(tlb_hit(comparator, ptr));
    }
# endif
    iotlbentry = &env_tlb(env)->d[ptr_mmu_idx].iotlb[index];

    /* If the virtual page MemAttr != Tagged, access unchecked. */
    if (!arm_tlb_mte_tagged(&iotlbentry->attrs)) {
        return NULL;
    }

    /*
     * If not backed by host ram, there is no tag storage: access unchecked.
     * This is probably a guest os bug though, so log it.
     */
    if (unlikely(flags & TLB_MMIO)) {
        qemu_log_mask(LOG_GUEST_ERROR,
                      "Page @ 0x%" PRIx64 " indicates Tagged Normal memory "
                      "but is not backed by host ram\n", ptr);
        return NULL;
    }

    /*
     * The Normal memory access can extend to the next page.  E.g. a single
     * 8-byte access to the last byte of a page will check only the last
     * tag on the first page.
     * Any page access exception has priority over tag check exception.
     */
    in_page = -(ptr | TARGET_PAGE_MASK);
    if (unlikely(ptr_size > in_page)) {
        void *ignore;
        flags |= probe_access_flags(env, ptr + in_page, ptr_access,
                                    ptr_mmu_idx, ra == 0, &ignore, ra);
        assert(!(flags & TLB_INVALID_MASK));
    }

    /* Any debug exception has priority over a tag check exception. */
    if (unlikely(flags & TLB_WATCHPOINT)) {
        int wp = ptr_access == MMU_DATA_LOAD ? BP_MEM_READ : BP_MEM_WRITE;
        assert(ra != 0);
        cpu_check_watchpoint(env_cpu(env), ptr, ptr_size,
                             iotlbentry->attrs, wp, ra);
    }

    

    /*
     * Find the physical address within the normal mem space.
     * The memory region lookup must succeed because TLB_MMIO was
     * not set in the cputlb lookup above.
     */
    mr = memory_region_from_host(host, &ptr_ra);
    tcg_debug_assert(mr != NULL);
    tcg_debug_assert(memory_region_is_ram(mr));
    ptr_paddr = ptr_ra;
    do {
        ptr_paddr += mr->addr;
        mr = mr->container;
    } while (mr);

    /* Convert to the physical address in tag space.  */
    tag_paddr = ptr_paddr >> (LOG2_TAG_GRANULE + 1);

    /* Look up the address in tag space. */
    tag_asi = iotlbentry->attrs.secure ? ARMASIdx_TagS : ARMASIdx_TagNS;
    tag_as = cpu_get_address_space(env_cpu(env), tag_asi);
    mr = address_space_translate(tag_as, tag_paddr, &xlat, NULL,
                                 tag_access == MMU_DATA_STORE,
                                 iotlbentry->attrs);

    /*
     * Note that @mr will never be NULL.  If there is nothing in the address
     * space at @tag_paddr, the translation will return the unallocated memory
     * region.  For our purposes, the result must be ram.
     */
    if (unlikely(!memory_region_is_ram(mr))) {
        /* ??? Failure is a board configuration error. */
        qemu_log_mask(LOG_UNIMP,
                      "Tag Memory @ 0x%" HWADDR_PRIx " not found for "
                      "Normal Memory @ 0x%" HWADDR_PRIx "\n",
                      tag_paddr, ptr_paddr);
        return NULL;
    }

    /*
     * Ensure the tag memory is dirty on write, for migration.
     * Tag memory can never contain code or display memory (vga).
     */
    if (tag_access == MMU_DATA_STORE) {
        ram_addr_t tag_ra = memory_region_get_ram_addr(mr) + xlat;
        cpu_physical_memory_set_dirty_flag(tag_ra, DIRTY_MEMORY_MIGRATION);
    }

    // xlat: type hwaddr, which is uint64_t.
    // printf("####QEMU_In_alloc_tag_mem input ptr: %ld | ptr_ra: %ld | mr: %p | mrRam: %p | xlat:%li | (ptr_pa>>4): %li !!!\n", ptr, ptr_ra, mr, memory_region_get_ram_ptr(mr), xlat, ptr_paddr >> 4);

    // v1: return btree_map[memory_region_get_ram_ptr(mr)].get(xlat)
    //     if memory_region_get_ram_ptr(mr) not in map, initialize it.
    // v2: One BTree for entire RAM.
    // tag is getTag(root, ptr_PA>>4).
    if (!global_btree_buffer)
    {
        // 4MB reserved, similar to tag_mem in virt.c : RAMsz/32, where RAMsz=128MB
        global_btree_buffer = malloc(4*1024*1024);
        global_btree_buffer_first_avail = 0;
        last_printed_global_btree_buf_first_avail = 0;
        last_printed_total_alloc_freed_nodes = 0;
        initBTree(global_btree_buffer);

        global_stg_call_count = 0;
    }
    if (bt_tag_addr)
    {
        uint64_t actual_tag_Addr = 2*xlat + ((ptr_paddr >> 4)%2);
        uint64_t max_allowd_addr = (uint64_t)1 << 32;
        if (actual_tag_Addr > max_allowd_addr)
            printf("BTREE_QEMU: ERRORRRRR!!!! Tag addr %li is more than 2^32!!! \n", actual_tag_Addr);
        *bt_tag_addr = (int)actual_tag_Addr;
    }

    return memory_region_get_ram_ptr(mr) + xlat; // Assuming mr = tag_sysmem, xlat has to be < 4MB.
#endif
}

uint64_t HELPER(irg)(CPUARMState *env, uint64_t rn, uint64_t rm)
{
    uint16_t exclude = extract32(rm | env->cp15.gcr_el1, 0, 16);
    int rrnd = extract32(env->cp15.gcr_el1, 16, 1);
    int start = extract32(env->cp15.rgsr_el1, 0, 4);
    int seed = extract32(env->cp15.rgsr_el1, 8, 16);
    int offset, i, rtag;

    /*
     * Our IMPDEF choice for GCR_EL1.RRND==1 is to continue to use the
     * deterministic algorithm.  Except that with RRND==1 the kernel is
     * not required to have set RGSR_EL1.SEED != 0, which is required for
     * the deterministic algorithm to function.  So we force a non-zero
     * SEED for that case.
     */
    if (unlikely(seed == 0) && rrnd) {
        do {
            Error *err = NULL;
            uint16_t two;

            if (qemu_guest_getrandom(&two, sizeof(two), &err) < 0) {
                /*
                 * Failed, for unknown reasons in the crypto subsystem.
                 * Best we can do is log the reason and use a constant seed.
                 */
                qemu_log_mask(LOG_UNIMP, "IRG: Crypto failure: %s\n",
                              error_get_pretty(err));
                error_free(err);
                two = 1;
            }
            seed = two;
        } while (seed == 0);
    }

    /* RandomTag */
    for (i = offset = 0; i < 4; ++i) {
        /* NextRandomTagBit */
        int top = (extract32(seed, 5, 1) ^ extract32(seed, 3, 1) ^
                   extract32(seed, 2, 1) ^ extract32(seed, 0, 1));
        seed = (top << 15) | (seed >> 1);
        offset |= top << i;
    }
    rtag = choose_nonexcluded_tag(start, offset, exclude);
    env->cp15.rgsr_el1 = rtag | (seed << 8);

    return address_with_allocation_tag(rn, rtag);
}

uint64_t HELPER(addsubg)(CPUARMState *env, uint64_t ptr,
                         int32_t offset, uint32_t tag_offset)
{
    int start_tag = allocation_tag_from_addr(ptr);
    uint16_t exclude = extract32(env->cp15.gcr_el1, 0, 16);
    int rtag = choose_nonexcluded_tag(start_tag, tag_offset, exclude);

    return address_with_allocation_tag(ptr + offset, rtag);
}

static int load_tag1(uint64_t ptr, uint8_t *mem)
{
    int ofs = extract32(ptr, LOG2_TAG_GRANULE, 1) * 4;
    return extract32(*mem, ofs, 4);
}

uint64_t HELPER(ldg)(CPUARMState *env, uint64_t ptr, uint64_t xt)
{
    int mmu_idx = cpu_mmu_index(env, false);
    uint8_t *mem;
    int rtag = 0, rtag2 = 0;

    /* Trap if accessing an invalid page.  */
    // v2 BTree: passing ptr to hwaddr
    int* bt_tag_addr = malloc(sizeof(int));
    mem = allocation_tag_mem(env, mmu_idx, ptr, MMU_DATA_LOAD, 1,
                             MMU_DATA_LOAD, 1, GETPC(), bt_tag_addr);

    /* Load if page supports tags. */
    if (mem) {
        rtag = load_tag1(ptr, mem);
        // printf("####QEMU:In_ldg_ptr:%ld | xt:%li | rtag:%i\n", ptr, xt, rtag);
        // v2 BTree: 
        rtag2 = getBTreeTag(getGlobalRootPtr(), *bt_tag_addr);
        // printf("$$$$BTREE_QEMU: ldg: Got tag_addr value from alloc_tag_mem: stored at: %p , value: %i, got tag: %i (rtag: %i), first_avail: %lli | root_offset: %i \n", bt_tag_addr, *bt_tag_addr, rtag2, rtag, global_btree_buffer_first_avail, global_btree_root_offset);
        // printf("$$$$BTREE_QEMU: current global_btree_buffer_first_avail: %lli, root_offset: %i \n", global_btree_buffer_first_avail, global_btree_root_offset);
        
        if (rtag != rtag2)
            printf("$$$$BTREE_QEMU: ldg: ERRORRRRRR: bt_tag %i != mem_tag %i \n", rtag2, rtag);
    }

    return address_with_allocation_tag(xt, rtag);
}

static void check_tag_aligned(CPUARMState *env, uint64_t ptr, uintptr_t ra)
{
    if (unlikely(!QEMU_IS_ALIGNED(ptr, TAG_GRANULE))) {
        // printf("ABOUT TO call arm_cpu_do_unaligned_access!!!!! \n");
        arm_cpu_do_unaligned_access(env_cpu(env), ptr, MMU_DATA_STORE,
                                    cpu_mmu_index(env, false), ra);
        g_assert_not_reached();
    }
}

/* For use in a non-parallel context, store to the given nibble.  */
static void store_tag1(uint64_t ptr, uint8_t *mem, int tag)
{
    int ofs = extract32(ptr, LOG2_TAG_GRANULE, 1) * 4;
    *mem = deposit32(*mem, ofs, 4, tag);
}

/* For use in a parallel context, atomically store to the given nibble.  */
static void store_tag1_parallel(uint64_t ptr, uint8_t *mem, int tag)
{
    int ofs = extract32(ptr, LOG2_TAG_GRANULE, 1) * 4;
    uint8_t old = qatomic_read(mem);

    while (1) {
        uint8_t new = deposit32(old, ofs, 4, tag);
        uint8_t cmp = qatomic_cmpxchg(mem, old, new);
        if (likely(cmp == old)) {
            return;
        }
        old = cmp;
    }
}

typedef void stg_store1(uint64_t, uint8_t *, int);

static inline void do_stg(CPUARMState *env, uint64_t ptr, uint64_t xt,
                          uintptr_t ra, stg_store1 store1)
{
    // printf("####QEMU:In_do_stg_ptr:%ld | xt:%li \n", ptr, xt);    
    int mmu_idx = cpu_mmu_index(env, false);
    uint8_t *mem;

    check_tag_aligned(env, ptr, ra);

    /* Trap if accessing an invalid page.  */
    int* bt_tag_addr = malloc(sizeof(int));
    mem = allocation_tag_mem(env, mmu_idx, ptr, MMU_DATA_STORE, TAG_GRANULE,
                             MMU_DATA_STORE, 1, ra, bt_tag_addr);

    /* Store if page supports tags. */
    if (mem) {
        int tag = allocation_tag_from_addr(xt);
        store1(ptr, mem, allocation_tag_from_addr(xt));
        setBTreeTag(getGlobalRootPtr(), *bt_tag_addr, tag);

        global_stg_call_count += 1;
        // 2000 for apache, 2000000, 100000000 for mysqld, || ((global_stg_call_count > 2000000) && (global_stg_call_count%2000000 == 0)) 
        if ((global_stg_call_count % 100000000 == 0) )
        {
            printf("SETTING Tag = %i for addr = %i (VA: %li) (stg_ct: %lli), global_btree_buffer_first_avail: %lli, alloc_freed_nodes: %i, total_alloc_freed_nodes: %i \n", tag, *bt_tag_addr, ptr, global_stg_call_count, global_btree_buffer_first_avail, alloc_freed_nodes, total_alloc_freed_nodes);
            last_printed_global_btree_buf_first_avail = global_btree_buffer_first_avail;
            last_printed_total_alloc_freed_nodes = total_alloc_freed_nodes;
        }
           
        if ( ((global_btree_buffer_first_avail > last_printed_global_btree_buf_first_avail) || (total_alloc_freed_nodes > last_printed_total_alloc_freed_nodes) ) && (last_printed_global_btree_buf_first_avail > 1900000)) // 1900000 for mysqld, 3900 for apache
        {
          last_printed_global_btree_buf_first_avail = global_btree_buffer_first_avail;
          last_printed_total_alloc_freed_nodes = total_alloc_freed_nodes;
          printf("IN STG: UPDATED global_btree_buffer_first_avail: %lli, alloc_freed_nodes: %i, total_alloc_freed_nodes: %i \n", global_btree_buffer_first_avail, alloc_freed_nodes, total_alloc_freed_nodes);
        }

        // int curr_tag = getBTreeTag(getGlobalRootPtr(), *bt_tag_addr); 
        // if (tag != 0) && (curr_tag != 0) && (tag != curr_tag)
        // {

        // }

        // print only at first verify failure:
        if ( (!btree_verify_failed) && !verifyBTreePtrs(getGlobalRootPtr()))
        {
            printf("ERRORRRRRR WHILE SETTING TAG = %i for addr = %i (stg_ct: %lli) \n", tag, *bt_tag_addr, global_stg_call_count);
            printBTree(getGlobalRootPtr(), 0);
            btree_verify_failed = true;
        }
        // else if ((global_stg_call_count > 5700) && (global_stg_call_count % 100 == 0)) // first tree prints around malloc-55 for apache19351.
        //     printBTree(getGlobalRootPtr(), 0);

        // printf("$$$$BTREE_QEMU: current global_btree_buffer_first_avail: %lli, root_offset: %i, alloc_freed_nodes: %i \n", global_btree_buffer_first_avail, global_btree_root_offset, alloc_freed_nodes);
        // , all_time_max_aval: %lli
    }
}

void HELPER(stg)(CPUARMState *env, uint64_t ptr, uint64_t xt)
{
    do_stg(env, ptr, xt, GETPC(), store_tag1);
}

void HELPER(stg_parallel)(CPUARMState *env, uint64_t ptr, uint64_t xt)
{
    do_stg(env, ptr, xt, GETPC(), store_tag1_parallel);
}

void HELPER(stg_stub)(CPUARMState *env, uint64_t ptr)
{
    int mmu_idx = cpu_mmu_index(env, false);
    uintptr_t ra = GETPC();

    check_tag_aligned(env, ptr, ra);
    probe_write(env, ptr, TAG_GRANULE, mmu_idx, ra);
}

static inline void do_st2g(CPUARMState *env, uint64_t ptr, uint64_t xt,
                           uintptr_t ra, stg_store1 store1)
{
    int mmu_idx = cpu_mmu_index(env, false);
    int tag = allocation_tag_from_addr(xt);
    uint8_t *mem1, *mem2;

    check_tag_aligned(env, ptr, ra);

    // printf("####QEMU:In_do_st2g_ptr:%ld | xt:%li | tag:%i\n", ptr, xt, tag);

    /*
     * Trap if accessing an invalid page(s).
     * This takes priority over !allocation_tag_access_enabled.
     */
    if (ptr & TAG_GRANULE) {
        /* Two stores unaligned mod TAG_GRANULE*2 -- modify two bytes. */
        mem1 = allocation_tag_mem(env, mmu_idx, ptr, MMU_DATA_STORE,
                                  TAG_GRANULE, MMU_DATA_STORE, 1, ra, NULL);
        mem2 = allocation_tag_mem(env, mmu_idx, ptr + TAG_GRANULE,
                                  MMU_DATA_STORE, TAG_GRANULE,
                                  MMU_DATA_STORE, 1, ra, NULL);

        /* Store if page(s) support tags. */
        if (mem1) {
            store1(TAG_GRANULE, mem1, tag);
        }
        if (mem2) {
            store1(0, mem2, tag);
        }
    } else {
        /* Two stores aligned mod TAG_GRANULE*2 -- modify one byte. */
        mem1 = allocation_tag_mem(env, mmu_idx, ptr, MMU_DATA_STORE,
                                  2 * TAG_GRANULE, MMU_DATA_STORE, 1, ra, NULL);
        if (mem1) {
            tag |= tag << 4; // Aditi: tag now looks like 0.... tag tag. [tag repeated]
            qatomic_set(mem1, tag);
        }
    }
}

void HELPER(st2g)(CPUARMState *env, uint64_t ptr, uint64_t xt)
{
    do_st2g(env, ptr, xt, GETPC(), store_tag1);
}

void HELPER(st2g_parallel)(CPUARMState *env, uint64_t ptr, uint64_t xt)
{
    do_st2g(env, ptr, xt, GETPC(), store_tag1_parallel);
}

void HELPER(st2g_stub)(CPUARMState *env, uint64_t ptr)
{
    int mmu_idx = cpu_mmu_index(env, false);
    uintptr_t ra = GETPC();
    int in_page = -(ptr | TARGET_PAGE_MASK);

    check_tag_aligned(env, ptr, ra);

    if (likely(in_page >= 2 * TAG_GRANULE)) {
        probe_write(env, ptr, 2 * TAG_GRANULE, mmu_idx, ra);
    } else {
        probe_write(env, ptr, TAG_GRANULE, mmu_idx, ra);
        probe_write(env, ptr + TAG_GRANULE, TAG_GRANULE, mmu_idx, ra);
    }
}

#define LDGM_STGM_SIZE  (4 << GMID_EL1_BS)

uint64_t HELPER(ldgm)(CPUARMState *env, uint64_t ptr)
{
    int mmu_idx = cpu_mmu_index(env, false);
    uintptr_t ra = GETPC();
    void *tag_mem;

    ptr = QEMU_ALIGN_DOWN(ptr, LDGM_STGM_SIZE);

    /* Trap if accessing an invalid page.  */
    tag_mem = allocation_tag_mem(env, mmu_idx, ptr, MMU_DATA_LOAD,
                                 LDGM_STGM_SIZE, MMU_DATA_LOAD,
                                 LDGM_STGM_SIZE / (2 * TAG_GRANULE), ra, NULL);

    /* The tag is squashed to zero if the page does not support tags.  */
    if (!tag_mem) {
        return 0;
    }

    // printf("####QEMU:In_ldgm_ptr:%ld_tag_mem:%p\n", ptr, tag_mem);

    QEMU_BUILD_BUG_ON(GMID_EL1_BS != 6);
    /*
     * We are loading 64-bits worth of tags.  The ordering of elements
     * within the word corresponds to a 64-bit little-endian operation.
     */
    return ldq_le_p(tag_mem);
}

void HELPER(stgm)(CPUARMState *env, uint64_t ptr, uint64_t val)
{
    int mmu_idx = cpu_mmu_index(env, false);
    uintptr_t ra = GETPC();
    void *tag_mem;

    ptr = QEMU_ALIGN_DOWN(ptr, LDGM_STGM_SIZE);

    /* Trap if accessing an invalid page.  */
    tag_mem = allocation_tag_mem(env, mmu_idx, ptr, MMU_DATA_STORE,
                                 LDGM_STGM_SIZE, MMU_DATA_LOAD,
                                 LDGM_STGM_SIZE / (2 * TAG_GRANULE), ra, NULL);

    /*
     * Tag store only happens if the page support tags,
     * and if the OS has enabled access to the tags.
     */
    if (!tag_mem) {
        return;
    }

    // printf("####QEMU:In_stgm_ptr:%ld | (16tags)val:%li | tag_mem:%p\n", ptr, val, tag_mem);

    QEMU_BUILD_BUG_ON(GMID_EL1_BS != 6);
    /*
     * We are storing 64-bits worth of tags.  The ordering of elements
     * within the word corresponds to a 64-bit little-endian operation.
     */
    stq_le_p(tag_mem, val);
}

void HELPER(stzgm_tags)(CPUARMState *env, uint64_t ptr, uint64_t val)
{
    uintptr_t ra = GETPC();
    int mmu_idx = cpu_mmu_index(env, false);
    int log2_dcz_bytes, log2_tag_bytes;
    intptr_t dcz_bytes, tag_bytes;
    uint8_t *mem;

    /*
     * In arm_cpu_realizefn, we assert that dcz > LOG2_TAG_GRANULE+1,
     * i.e. 32 bytes, which is an unreasonably small dcz anyway,
     * to make sure that we can access one complete tag byte here.
     */
    log2_dcz_bytes = env_archcpu(env)->dcz_blocksize + 2;
    log2_tag_bytes = log2_dcz_bytes - (LOG2_TAG_GRANULE + 1);
    dcz_bytes = (intptr_t)1 << log2_dcz_bytes;
    tag_bytes = (intptr_t)1 << log2_tag_bytes;
    ptr &= -dcz_bytes;

    mem = allocation_tag_mem(env, mmu_idx, ptr, MMU_DATA_STORE, dcz_bytes,
                             MMU_DATA_STORE, tag_bytes, ra, NULL);
    if (mem) {
        // printf("####QEMU:In_stzgm_tags_ptr:%ldval:%li\n", ptr, val);
        int tag_pair = (val & 0xf) * 0x11;
        memset(mem, tag_pair, tag_bytes);
    }
}

/* Record a tag check failure.  */
static void mte_check_fail(CPUARMState *env, uint32_t desc,
                           uint64_t dirty_ptr, uintptr_t ra)
{
    int mmu_idx = FIELD_EX32(desc, MTEDESC, MIDX);
    ARMMMUIdx arm_mmu_idx = core_to_aa64_mmu_idx(mmu_idx);
    int el, reg_el, tcf, select, is_write, syn;
    uint64_t sctlr;

    reg_el = regime_el(env, arm_mmu_idx);
    sctlr = env->cp15.sctlr_el[reg_el];

    switch (arm_mmu_idx) {
    case ARMMMUIdx_E10_0:
    case ARMMMUIdx_E20_0:
        el = 0;
        tcf = extract64(sctlr, 38, 2);
        break;
    default:
        el = reg_el;
        tcf = extract64(sctlr, 40, 2);
    }

    switch (tcf) {
    case 1:
        /*
         * Tag check fail causes a synchronous exception.
         *
         * In restore_state_to_opc, we set the exception syndrome
         * for the load or store operation.  Unwind first so we
         * may overwrite that with the syndrome for the tag check.
         */
        cpu_restore_state(env_cpu(env), ra, true);
        env->exception.vaddress = dirty_ptr;

        is_write = FIELD_EX32(desc, MTEDESC, WRITE);
        syn = syn_data_abort_no_iss(arm_current_el(env) != 0, 0, 0, 0, 0,
                                    is_write, 0x11);
        raise_exception(env, EXCP_DATA_ABORT, syn, exception_target_el(env));
        /* noreturn, but fall through to the assert anyway */

    case 0:
        /*
         * Tag check fail does not affect the PE.
         * We eliminate this case by not setting MTE_ACTIVE
         * in tb_flags, so that we never make this runtime call.
         */
        g_assert_not_reached();

    case 2:
        /* Tag check fail causes asynchronous flag set.  */
        if (regime_has_2_ranges(arm_mmu_idx)) {
            select = extract64(dirty_ptr, 55, 1);
        } else {
            select = 0;
        }
        env->cp15.tfsr_el[el] |= 1 << select;
#ifdef CONFIG_USER_ONLY
        /*
         * Stand in for a timer irq, setting _TIF_MTE_ASYNC_FAULT,
         * which then sends a SIGSEGV when the thread is next scheduled.
         * This cpu will return to the main loop at the end of the TB,
         * which is rather sooner than "normal".  But the alternative
         * is waiting until the next syscall.
         */
        qemu_cpu_kick(env_cpu(env));
#endif
        break;

    default:
        /* Case 3: Reserved. */
        qemu_log_mask(LOG_GUEST_ERROR,
                      "Tag check failure with SCTLR_EL%d.TCF%s "
                      "set to reserved value %d\n",
                      reg_el, el ? "" : "0", tcf);
        break;
    }
}

/*
 * Perform an MTE checked access for a single logical or atomic access.
 */
static bool mte_probe1_int(CPUARMState *env, uint32_t desc, uint64_t ptr,
                           uintptr_t ra, int bit55)
{
    int mem_tag, mmu_idx, ptr_tag, size, bt_mem_tag;
    MMUAccessType type;
    uint8_t *mem;

    ptr_tag = allocation_tag_from_addr(ptr);

    if (tcma_check(desc, bit55, ptr_tag)) {
        return true;
    }

    mmu_idx = FIELD_EX32(desc, MTEDESC, MIDX);
    type = FIELD_EX32(desc, MTEDESC, WRITE) ? MMU_DATA_STORE : MMU_DATA_LOAD;
    size = FIELD_EX32(desc, MTEDESC, ESIZE);

    int* bt_tag_addr = malloc(sizeof(int));
    mem = allocation_tag_mem(env, mmu_idx, ptr, type, size,
                             MMU_DATA_LOAD, 1, ra, bt_tag_addr);
    if (!mem) {
        return true;
    }

    mem_tag = load_tag1(ptr, mem);
    bt_mem_tag = getBTreeTag(getGlobalRootPtr(), *bt_tag_addr);
    // printf("$$$$BTREE_QEMU: mte_probe1_int: %ld | ra:%li | bt_tag_addr: %i | ptr_tag:%i | mem_tag:%i | tag from btree: %i | current first_avail: %lli | root_offset: %i | alloc_freed_nodes: %i \n", ptr, ra, *bt_tag_addr, ptr_tag, mem_tag, bt_mem_tag, global_btree_buffer_first_avail, global_btree_root_offset, alloc_freed_nodes);
    // // printf("$$$$BTREE_QEMU: mte_probe1_int: tag_addr: %i, mem_tag: %i, tag from btree: %i \n", *bt_tag_addr, mem_tag, bt_mem_tag);
    // // printf("$$$$BTREE_QEMU: current global_btree_buffer_first_avail: %lli, root_offset: %i \n", global_btree_buffer_first_avail, global_btree_root_offset);
    if (bt_mem_tag != mem_tag)
    {
        printf("$$$$BTREE_QEMU: mte_probe1_int: ERRORRRRRR!!! tag addr: %i, bt_mem_tag %i != mem_tag %i \n", *bt_tag_addr, bt_mem_tag, mem_tag);
        if (!btree_verify_failed)
        {
          printBTree(getGlobalRootPtr(), 0);
          btree_verify_failed = true;
        }
    }
    return ptr_tag == mem_tag;
}

/*
 * No-fault version of mte_check1, to be used by SVE for MemSingleNF.
 * Returns false if the access is Checked and the check failed.  This
 * is only intended to probe the tag -- the validity of the page must
 * be checked beforehand.
 */
bool mte_probe1(CPUARMState *env, uint32_t desc, uint64_t ptr)
{
    int bit55 = extract64(ptr, 55, 1);

    /* If TBI is disabled, the access is unchecked. */
    if (unlikely(!tbi_check(desc, bit55))) {
        return true;
    }

    return mte_probe1_int(env, desc, ptr, 0, bit55);
}

uint64_t mte_check1(CPUARMState *env, uint32_t desc,
                    uint64_t ptr, uintptr_t ra)
{
    int bit55 = extract64(ptr, 55, 1);

    /* If TBI is disabled, the access is unchecked, and ptr is not dirty. */
    if (unlikely(!tbi_check(desc, bit55))) {
        return ptr;
    }

    if (unlikely(!mte_probe1_int(env, desc, ptr, ra, bit55))) {
        mte_check_fail(env, desc, ptr, ra);
    }

    return useronly_clean_ptr(ptr);
}

uint64_t HELPER(mte_check1)(CPUARMState *env, uint32_t desc, uint64_t ptr)
{
    return mte_check1(env, desc, ptr, GETPC());
}

/*
 * Perform an MTE checked access for multiple logical accesses.
 */

/**
 * checkN:
 * @tag: tag memory to test
 * @odd: true to begin testing at tags at odd nibble
 * @cmp: the tag to compare against
 * @count: number of tags to test
 *
 * Return the number of successful tests.
 * Thus a return value < @count indicates a failure.
 *
 * A note about sizes: count is expected to be small.
 *
 * The most common use will be LDP/STP of two integer registers,
 * which means 16 bytes of memory touching at most 2 tags, but
 * often the access is aligned and thus just 1 tag.
 *
 * Using AdvSIMD LD/ST (multiple), one can access 64 bytes of memory,
 * touching at most 5 tags.  SVE LDR/STR (vector) with the default
 * vector length is also 64 bytes; the maximum architectural length
 * is 256 bytes touching at most 9 tags.
 *
 * The loop below uses 7 logical operations and 1 memory operation
 * per tag pair.  An implementation that loads an aligned word and
 * uses masking to ignore adjacent tags requires 18 logical operations
 * and thus does not begin to pay off until 6 tags.
 * Which, according to the survey above, is unlikely to be common.
 */
static int checkN(uint8_t *mem, int odd, int cmp, int count)
{
    int n = 0, diff;

    /* Replicate the test tag and compare.  */
    cmp *= 0x11;
    diff = *mem++ ^ cmp;

    if (odd) {
        goto start_odd;
    }

    while (1) {
        /* Test even tag. */
        if (unlikely((diff) & 0x0f)) {
            break;
        }
        if (++n == count) {
            break;
        }

    start_odd:
        /* Test odd tag. */
        if (unlikely((diff) & 0xf0)) {
            break;
        }
        if (++n == count) {
            break;
        }

        diff = *mem++ ^ cmp;
    }
    return n;
}

uint64_t mte_checkN(CPUARMState *env, uint32_t desc,
                    uint64_t ptr, uintptr_t ra)
{
    int mmu_idx, ptr_tag, bit55;
    uint64_t ptr_last, ptr_end, prev_page, next_page;
    uint64_t tag_first, tag_end;
    uint64_t tag_byte_first, tag_byte_end;
    uint32_t esize, total, tag_count, tag_size, n, c;
    uint8_t *mem1, *mem2;
    MMUAccessType type;

    bit55 = extract64(ptr, 55, 1);

    /* If TBI is disabled, the access is unchecked, and ptr is not dirty. */
    if (unlikely(!tbi_check(desc, bit55))) {
        return ptr;
    }

    ptr_tag = allocation_tag_from_addr(ptr);

    // printf("####QEMU:In_mte_checkN_ptr:%ld_ra:%li_ptr_tag:%i\n", ptr, ra, ptr_tag);

    if (tcma_check(desc, bit55, ptr_tag)) {
        goto done;
    }

    mmu_idx = FIELD_EX32(desc, MTEDESC, MIDX);
    type = FIELD_EX32(desc, MTEDESC, WRITE) ? MMU_DATA_STORE : MMU_DATA_LOAD;
    esize = FIELD_EX32(desc, MTEDESC, ESIZE);
    total = FIELD_EX32(desc, MTEDESC, TSIZE);

    /* Find the addr of the end of the access, and of the last element. */
    ptr_end = ptr + total;
    ptr_last = ptr_end - esize;

    /* Round the bounds to the tag granule, and compute the number of tags. */
    tag_first = QEMU_ALIGN_DOWN(ptr, TAG_GRANULE);
    tag_end = QEMU_ALIGN_UP(ptr_last, TAG_GRANULE);
    tag_count = (tag_end - tag_first) / TAG_GRANULE;

    /* Round the bounds to twice the tag granule, and compute the bytes. */
    tag_byte_first = QEMU_ALIGN_DOWN(ptr, 2 * TAG_GRANULE);
    tag_byte_end = QEMU_ALIGN_UP(ptr_last, 2 * TAG_GRANULE);

    /* Locate the page boundaries. */
    prev_page = ptr & TARGET_PAGE_MASK;
    next_page = prev_page + TARGET_PAGE_SIZE;

    if (likely(tag_end - prev_page <= TARGET_PAGE_SIZE)) {
        /* Memory access stays on one page. */
        tag_size = (tag_byte_end - tag_byte_first) / (2 * TAG_GRANULE);
        mem1 = allocation_tag_mem(env, mmu_idx, ptr, type, total,
                                  MMU_DATA_LOAD, tag_size, ra, NULL);
        if (!mem1) {
            goto done;
        }
        /* Perform all of the comparisons. */
        n = checkN(mem1, ptr & TAG_GRANULE, ptr_tag, tag_count);
    } else {
        /* Memory access crosses to next page. */
        tag_size = (next_page - tag_byte_first) / (2 * TAG_GRANULE);
        mem1 = allocation_tag_mem(env, mmu_idx, ptr, type, next_page - ptr,
                                  MMU_DATA_LOAD, tag_size, ra, NULL);

        tag_size = (tag_byte_end - next_page) / (2 * TAG_GRANULE);
        mem2 = allocation_tag_mem(env, mmu_idx, next_page, type,
                                  ptr_end - next_page,
                                  MMU_DATA_LOAD, tag_size, ra, NULL);

        /*
         * Perform all of the comparisons.
         * Note the possible but unlikely case of the operation spanning
         * two pages that do not both have tagging enabled.
         */
        n = c = (next_page - tag_first) / TAG_GRANULE;
        if (mem1) {
            n = checkN(mem1, ptr & TAG_GRANULE, ptr_tag, c);
        }
        if (n == c) {
            if (!mem2) {
                goto done;
            }
            n += checkN(mem2, 0, ptr_tag, tag_count - c);
        }
    }

    /*
     * If we failed, we know which granule.  Compute the element that
     * is first in that granule, and signal failure on that element.
     */
    if (unlikely(n < tag_count)) {
        uint64_t fail_ofs;

        fail_ofs = tag_first + n * TAG_GRANULE - ptr;
        fail_ofs = ROUND_UP(fail_ofs, esize);
        mte_check_fail(env, desc, ptr + fail_ofs, ra);
    }

 done:
    return useronly_clean_ptr(ptr);
}

uint64_t HELPER(mte_checkN)(CPUARMState *env, uint32_t desc, uint64_t ptr)
{
    return mte_checkN(env, desc, ptr, GETPC());
}

/*
 * Perform an MTE checked access for DC_ZVA.
 */
uint64_t HELPER(mte_check_zva)(CPUARMState *env, uint32_t desc, uint64_t ptr)
{
    uintptr_t ra = GETPC();
    int log2_dcz_bytes, log2_tag_bytes;
    int mmu_idx, bit55;
    intptr_t dcz_bytes, tag_bytes, i;
    void *mem;
    uint64_t ptr_tag, mem_tag, align_ptr;

    bit55 = extract64(ptr, 55, 1);

    /* If TBI is disabled, the access is unchecked, and ptr is not dirty. */
    if (unlikely(!tbi_check(desc, bit55))) {
        return ptr;
    }

    ptr_tag = allocation_tag_from_addr(ptr);

    if (tcma_check(desc, bit55, ptr_tag)) {
        goto done;
    }

    // printf("####QEMU:In_mte_check_zva_ptr:%ld_ptr_tag:%li\n", ptr, ptr_tag);

    /*
     * In arm_cpu_realizefn, we asserted that dcz > LOG2_TAG_GRANULE+1,
     * i.e. 32 bytes, which is an unreasonably small dcz anyway, to make
     * sure that we can access one complete tag byte here.
     */
    log2_dcz_bytes = env_archcpu(env)->dcz_blocksize + 2;
    log2_tag_bytes = log2_dcz_bytes - (LOG2_TAG_GRANULE + 1);
    dcz_bytes = (intptr_t)1 << log2_dcz_bytes;
    tag_bytes = (intptr_t)1 << log2_tag_bytes;
    align_ptr = ptr & -dcz_bytes;

    /*
     * Trap if accessing an invalid page.  DC_ZVA requires that we supply
     * the original pointer for an invalid page.  But watchpoints require
     * that we probe the actual space.  So do both.
     */
    mmu_idx = FIELD_EX32(desc, MTEDESC, MIDX);
    (void) probe_write(env, ptr, 1, mmu_idx, ra);
    mem = allocation_tag_mem(env, mmu_idx, align_ptr, MMU_DATA_STORE,
                             dcz_bytes, MMU_DATA_LOAD, tag_bytes, ra, NULL);
    if (!mem) {
        goto done;
    }



    /*
     * Unlike the reasoning for checkN, DC_ZVA is always aligned, and thus
     * it is quite easy to perform all of the comparisons at once without
     * any extra masking.
     *
     * The most common zva block size is 64; some of the thunderx cpus use
     * a block size of 128.  For user-only, aarch64_max_initfn will set the
     * block size to 512.  Fill out the other cases for future-proofing.
     *
     * In order to be able to find the first miscompare later, we want the
     * tag bytes to be in little-endian order.
     */
    switch (log2_tag_bytes) {
    case 0: /* zva_blocksize 32 */
        mem_tag = *(uint8_t *)mem;
        ptr_tag *= 0x11u;
        break;
    case 1: /* zva_blocksize 64 */
        mem_tag = cpu_to_le16(*(uint16_t *)mem);
        ptr_tag *= 0x1111u;
        break;
    case 2: /* zva_blocksize 128 */
        mem_tag = cpu_to_le32(*(uint32_t *)mem);
        ptr_tag *= 0x11111111u;
        break;
    case 3: /* zva_blocksize 256 */
        mem_tag = cpu_to_le64(*(uint64_t *)mem);
        ptr_tag *= 0x1111111111111111ull;
        break;

    default: /* zva_blocksize 512, 1024, 2048 */
        ptr_tag *= 0x1111111111111111ull;
        i = 0;
        do {
            mem_tag = cpu_to_le64(*(uint64_t *)(mem + i));
            if (unlikely(mem_tag != ptr_tag)) {
                goto fail;
            }
            i += 8;
            align_ptr += 16 * TAG_GRANULE;
        } while (i < tag_bytes);
        goto done;
    }

    if (likely(mem_tag == ptr_tag)) {
        goto done;
    }

 fail:
    /* Locate the first nibble that differs. */
    i = ctz64(mem_tag ^ ptr_tag) >> 4;
    mte_check_fail(env, desc, align_ptr + i * TAG_GRANULE, ra);

 done:
    return useronly_clean_ptr(ptr);
}