BufferTree.java

package avl;

import java.io.File;
import java.io.Closeable;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.nio.channels.FileChannel;
import java.nio.ByteBuffer;
import java.nio.IntBuffer;
import java.util.List;
import java.util.LinkedList;

public class BufferTree implements Closeable {

  private static final int LEFT = -1;
  private static final int RIGHT = 1;

  public static final int ELEMENT_SIZE = 4 * Integer.BYTES;
  public static final int META_SIZE = 2 * Integer.BYTES;
  public static final int NULL = 0;

  // META offsets
  private static final int NEXT_AVAILABLE_OFFSET = 0;
  private static final int ROOT_OFFSET = 1;

  // Node offsets
  private static final int VALUE_OFFSET = 0;
  private static final int LEFT_OFFSET = 1;
  private static final int RIGHT_OFFSET = 2;
  private static final int BALANCE_OFFSET = 3;

  private final RandomAccessFile file;
  private final FileChannel fileChannel;
  private ByteBuffer buffer;
  private IntBuffer metaBuffer;
  private Node root;

  public BufferTree() throws IOException {
    this("buffer.bin", 20);
  }

  public BufferTree(int length) throws IOException {
    this("buffer.bin", length);
  }

  public BufferTree(String filename) throws IOException {
    this(filename, 20);
  }

  public BufferTree(String filename, int length) throws IOException {
    long fileLength = ELEMENT_SIZE * length;
    File ioFile = new File(filename);
    boolean exists = ioFile.exists();
    file = new RandomAccessFile(filename, "rw");
    if (exists && file.length() > fileLength) {
      fileLength = file.length();
    } else if (fileLength > file.length()) {
      file.setLength(fileLength);
    }

    fileChannel = file.getChannel();
    buffer = fileChannel.map(FileChannel.MapMode.READ_WRITE, 0, fileLength);
    metaBuffer = buffer.limit(META_SIZE).position(0).slice().asIntBuffer();
    if (!exists) {
      setNextAvailable(1);
    }
    int rootIndex = metaBuffer.get(ROOT_OFFSET);
    root = rootIndex == NULL ? null : new Node(rootIndex);
  }

  public void close() throws IOException {
    fileChannel.force(true);
    file.close();
  }

  public Node add(int value) {
    Node node = new Node(getAndIncNextAvailable(), value);
    if (root == null) {
      setRoot(node);
    } else {
      Node newRoot = insertNode(root, node);
      if (root != newRoot) {
        setRoot(newRoot);
      }
    }
    return root;
  }

  private static List<Integer> appendToList(List<Integer> acc, Node node) {
    if (node == null) {
      return acc;
    } else {
      appendToList(acc, node.getChild(LEFT));
      acc.add(node.getValue());
      return appendToList(acc, node.getChild(RIGHT));
    }
  }

  public List<Integer> asList() {
    return appendToList(new LinkedList<Integer>(), root);
  }

  private void setNextAvailable(int next) {
    metaBuffer.put(NEXT_AVAILABLE_OFFSET, next);
  }

  private int getAndIncNextAvailable() {
    int next = metaBuffer.get(NEXT_AVAILABLE_OFFSET);
    if ((next + 1) * ELEMENT_SIZE > buffer.capacity()) {
      throw new RuntimeException("Out of capacity");
    }
    metaBuffer.put(NEXT_AVAILABLE_OFFSET, next + 1);
    return next;
  }

  private void setRoot(Node node) {
    root = node;
    metaBuffer.put(ROOT_OFFSET, node.getIndex());
  }

  public Node getRoot() {
    return root;
  }

  public String toString() {
    return root.toString();
  }

  private static final int other(int side) {
    return side == LEFT ? RIGHT : LEFT;
  }

  private static Node insertNode(Node node, Node newNode) {
    int value = node.getValue();

    int side = newNode.getValue() < value ? LEFT : RIGHT;

    if (node.getChild(side) == null) {
      return node.initChild(side, newNode);
    }

    int childBalance = node.getChild(side).getBalance();
    node.setChild(side, insertNode(node.getChild(side), newNode));

    // Did the child balance change from 0? Then it's deeper
    int balance = node.getBalance();
    if (childBalance == 0 && node.getChild(side).getBalance() != 0) {
      balance += side;
      node.setBalance(balance);
    }

    if (Math.abs(balance) == 2) {
      return rebalance(node, balance);
    }
    return node;
  }

  private static Node rebalance(Node node, int balance) {
    var side = balance < 0 ? LEFT : RIGHT;

    if (side == node.getChild(side).getBalance()) {
      return rebalanceSS(node, side);
    } else {
      return rebalanceSO(node, side);
    }
  }

  private static Node rebalanceSS(Node node, int side) {
    var nodeS = node.getChild(side);
    node.setChild(side, nodeS.getChild(other(side)));
    nodeS.setChild(other(side), node);
    node.setBalance(0);
    nodeS.setBalance(0);
    return nodeS;
  }

  private static Node rebalanceSO(Node node, int side) {
    var otherSide = other(side);

    var nodeS = node.getChild(side);
    var nodeSO = nodeS.getChild(otherSide);
    node.setChild(side, nodeSO.getChild(otherSide));
    nodeS.setChild(otherSide, nodeSO.getChild(side));
    nodeSO.setChild(otherSide, node);
    nodeSO.setChild(side, nodeS);
    if (nodeSO.getBalance() == otherSide) {
      node.setBalance(0);
      nodeS.setBalance(side);
    } else if (nodeSO.getBalance() == side) {
      node.setBalance(otherSide);
      nodeS.setBalance(0);
    } else {
      node.setBalance(0);
      nodeS.setBalance(0);
    }
    nodeSO.setBalance(0);
    return nodeSO;
  }

  public static String treeString(Node element) {
    Node left = element.getChild(LEFT);
    Node right = element.getChild(RIGHT);
    return (left == null ? "" : treeString(left)) +
           Integer.toString(element.getValue()) +
           (right == null ? "" : treeString(right));
  }

  public class Node {
    private final IntBuffer intBuffer;
    private final int index;

    Node(int index) {
      this.index = index;
      int offset = index * ELEMENT_SIZE;
      if (offset > buffer.limit()) {
        intBuffer = buffer.limit(offset + ELEMENT_SIZE).position(offset).slice().asIntBuffer();
      } else {
        intBuffer = buffer.position(offset).limit(offset + ELEMENT_SIZE).slice().asIntBuffer();
      }
    }

    Node(int index, int value) {
      this(index);
      setValue(value);
      setChild(LEFT, null);
      setChild(RIGHT, null);
      setBalance(0);
    }

    public int getIndex() {
      return index;
    }

    public int getValue() {
      return intBuffer.get(VALUE_OFFSET);
    }

    public Node getChild(int side) {
      int childId = intBuffer.get(side == LEFT ? LEFT_OFFSET : RIGHT_OFFSET);
      return childId == NULL ? null : new Node(childId);
    }

    public int getBalance() {
      return intBuffer.get(BALANCE_OFFSET);
    }

    public Node setValue(int value) {
      intBuffer.put(VALUE_OFFSET, value);
      return this;
    }

    public Node setChild(int side, Node child) {
      intBuffer.put(side == LEFT ? LEFT_OFFSET : RIGHT_OFFSET,
                    child == null ? NULL : child.getIndex());
      return this;
    }

    public Node initChild(int side, Node child) {
      setChild(side, child);
      setBalance(getBalance() + side);
      return this;
    }

    public Node setBalance(int balance) {
      if (Math.abs(balance) > 2) {
        throw new RuntimeException("balance out of range: " + balance);
      }
      intBuffer.put(BALANCE_OFFSET, balance);
      return this;
    }

    public String toString() {
      Node left = getChild(LEFT);
      Node right = getChild(RIGHT);
      return (left == null ? "" : left.toString() + ", ") +
             Integer.toString(getValue()) +
             (right == null ? "" : ", " + right.toString());
    }
  }
}

TreeExample.java

package avl;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Comparator;

public class TreeExample {

  static final int[] digits =
    new int[]
    {3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4, 6, 2,
     6, 4, 3, 3, 8, 3, 2, 7, 9, 5, 0, 2, 8, 8, 4, 1, 9, 7, 1, 6, 9, 3,
     9, 9, 3, 7, 5, 1, 0, 5, 8, 2, 0, 9, 7, 4, 9, 4, 4, 5, 9, 2, 3, 0};

  public static void main(String[] args) throws IOException {
    BufferTree bufferTree = new BufferTree("tree.bin", 100);
    for (var i: digits) {
      bufferTree.add(i);
    }
    System.out.println(bufferTree);

    // put the original digits in a list and sort them
    var digitList = new ArrayList<Integer>(digits.length);
    for (var i: digits) digitList.add(i);
    digitList.sort(Comparator.comparingInt(x->x));

    // compare the list from the tree to the sorted list
    if (bufferTree.asList().equals(digitList)) {
      System.out.println("Digits sorted correctly");
    }
    bufferTree.close();
  }
}

TreeExample2.java

package avl;

import java.io.IOException;

public class TreeExample2 {
  public static void main(String[] args) throws IOException {
    BufferTree bufferTree = new BufferTree("tree.bin", 25);
    System.out.println(bufferTree);
    bufferTree.close();
  }
}