// 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;
    });
  }
}