// A naive suffix tree implementation

sclass SuffixTree {

  Node root;

  S fullText;

  int nodeCount;

  Comparator<IFirstChar> childComparator = (a, b) -> a.firstCharOrMinus1(this)-b.firstCharOrMinus1(this);

  new DummyNode dummyNode;

  sinterface IFirstChar {

    int firstCharOrMinus1(SuffixTree tree);

  }

  sclass DummyNode implements IFirstChar {

    int firstChar;

    public int firstCharOrMinus1(SuffixTree tree) { ret firstChar; }

  }

  sclass Node implements IFirstChar {

    int from, to; // our range in fullText

    O children; // null, a Node or Node[] sorted by first character

    int position = -1;

    *() {}

    *(int *from, int *to) {}

    *(Substring s) { from = s.startIndex(); to = s.endIndex(); }

    Substring text(SuffixTree tree) { ret Substring(tree.fullText, from, to); }

    int lText() { ret to-from; }

    void addPosition(int i) {

      if (position >= 0) fail("Already have position");

      position = i;

    }

    void addChild(SuffixTree tree, Node n) {

      if (children == null) children = n;

      else if (children cast Node)

        children = sortArrayInPlace(new Node[] {children, n}, tree.childComparator);

      else {

        Node[] c = cast children;

        Node[] x = new[l(c)+1];

        arraycopy(c, 0, x, 0, l(c));

        x[l(x)-1] = n;

        children = sortArrayInPlace(x, tree.childComparator);

      }

    }

    Cl<Node> children() {

      if (children == null) ret emptyList();

      if (children cast Node) ret ll(children);

      ret wrapArrayAsList((Node[]) children);

    }

    CompactTreeSet<IFirstChar> makeEmptyChildrenSet(SuffixTree tree) { ret new CompactTreeSet<>(tree.childComparator); }

    public int firstCharOrMinus1(SuffixTree tree) {

      ret charToIntOrMinus1(text(tree));

    }

    Node getChild(SuffixTree tree, int c) {

      if (children == null) null;

      if (children cast Node)

        ret children.firstCharOrMinus1(tree) == c ? children : null;

      tree.dummyNode.firstChar = c;

      Node[] x = cast children;

      int i = Arrays.binarySearch(x, tree.dummyNode, tree.childComparator);

      ret i >= 0 ? x[i] : null;

    }

  }

  *() {}

  *(S *fullText) {

    root = new Node(0, 0);

    ++nodeCount;

    for i over fullText: {

      addSuffix(Substring(fullText, i), i);

      if (((i+1) % 1000000) == 0)

        print((i+1) + " suffixes added (" + nNodes(nodeCount) + ")");

    }

  }

  void addSuffix(Substring s, int position) {

    Node node = root;

    while (!empty(s)) {

      int _n = lCommonPrefix_CharSequence(node.text(SuffixTree.this), s);

      s = s.substring(_n);

      if (_n >= node.lText()) { // node text exhausted

        if (empty(s)) { // pattern also exhausted - done

          node.addPosition(position);

          ret;

        } else {

          Node n = node.getChild(SuffixTree.this, charToIntOrMinus1(s));

          if (n == null) {

            n = new Node(s);

            ++nodeCount;

            n.addPosition(position);

            node.addChild(SuffixTree.this, n);

            ret;

          } else

            node = n;

        }

      } else { // node text not exhausted

        // split node. first, move all the node's vitals to a new node nOld

        Node nOld = new Node(node.from+_n, node.to);

        ++nodeCount;

        nOld.position = node.position;

        node.position = -1;

        nOld.children = node.children;

        node.children = null;

        node.to = node.from+_n;

        node.addChild(SuffixTree.this, nOld);

        // now add a new node

        Node nNew = new Node(s);

        ++nodeCount;

        nNew.addPosition(position);

        node.addChild(SuffixTree.this, nNew);

        ret;

      }

    }

  }

  public L<Int> indicesOf(S pattern) {

    ret asList(indicesOf_iterator(pattern));

  }

  public ItIt<Int> indicesOf_iterator(S pattern) {

    ret mapI_notNull(allNodesUnder(scanDown(root, pattern)),

      n -> n.position >= 0 ? n.position : null);

  }

  Node scanDown(Node node, S pattern) {

    int lPattern = l(pattern), iPattern = 0;

    while true {

      int n = lCommonPrefix_CharSequence(node.text(this), Substring(pattern, iPattern));

      iPattern += n;

      if (iPattern >= lPattern) break; // pattern exhausted - done

      if (n < node.lText()) null; // mismatch, exit

      Node child = node.getChild(SuffixTree.this, charAtAsIntOrMinus1(pattern, iPattern));

      if (child != null) continue with node = child;

      null;

    }

    ret node;

  }

  L<Int> getPositions(Node node) {

    new IntBuffer out;

    collectPositions(out, node);

    ret out.asVirtualList();

  }

  void collectPositions(IntBuffer out, Node node) {

    if (node == null) ret;

    if (node.position >= 0) out.add(node.position);

    fOr (IFirstChar n : node.children())

      collectPositions(out, (Node) n);

  }

  void printMe() {

    printNode("", "", root);

  }

  void printNode(S indent, S pre, Node node) {

    print(indent + pre + quote(shorten(20, node.text(this))) + (node.position < 0 ? "" : " [" + node.position + "]"));

    fOr (IFirstChar _n : node.children()) {

      Node n = cast _n;

      printNode(indent + "  ", "[" + (n.lText() == 0 ? "end" : quote(n.text(this).charAt(0))) + "] ", n);

    }

  }

  ItIt<Node> allNodes() {

    ret allNodesUnder(root);

  }

  // includes the node itself

  ItIt<Node> allNodesUnder(Node node) {

    new L<Iterator<Node>> stack;

    if (node != null)

      stack.add(iteratorLL(node));

    ret iteratorFromFunction_if0(() -> {

      while (nempty(stack)) {

        if (!last(stack).hasNext())

          popLast(stack);

        else {

          Node n = last(stack).next();

          stack.add((Iterator) iterator(n.children()));

          ret n;

        }

      }

      null;

    });

  }

}