TreeTraversal: introduce INTERLEAVED_DFS

platform/platform-tests/testSrc/com/intellij/util/containers/TreeTraverserTest.java

@@ -15,6 +15,7 @@
  */
 package com.intellij.util.containers;
 
+import com.intellij.openapi.util.Comparing;
 import com.intellij.openapi.util.Condition;
 import com.intellij.util.Consumer;
 import com.intellij.util.Function;
@@ -146,6 +147,21 @@ public class TreeTraverserTest extends TestCase {
     };
   }
 
+  @NotNull
+  private static <E> JBIterable.StatefulFilter<E> UP_TO(final E o) {
+    return new JBIterable.StatefulFilter<E>() {
+      boolean b;
+
+      @Override
+      public boolean value(E e) {
+        if (b) return false;
+        b = Comparing.equal(e, o);
+        return true;
+      }
+    };
+  }
+
+
   // JBIterable ----------------------------------------------
 
   public void testAppend() {
@@ -248,6 +264,47 @@ public class TreeTraverserTest extends TestCase {
     assertEquals(Arrays.asList(1, 2, 5, 6, 7, 3, 8, 9, 10, 4, 11, 12, 13), numTraverser(TreeTraversal.PRE_ORDER_DFS).fun(1).toList());
   }
 
+  public void testSimpleInterlacedDfs() {
+    assertEquals(Arrays.asList(1, 2, 5, 3, 6, 4, 8, 7, 9, 11, 10, 12, 13), numTraverser(TreeTraversal.INTERLEAVED_DFS).fun(1).toList());
+  }
+
+  public void testCyclicInterlacedDfs() {
+    Function<Integer, JBIterable<Integer>> traversal = TreeTraversal.INTERLEAVED_DFS.traversal(Functions.fromMap(
+      ContainerUtil.<Integer, Collection<Integer>>immutableMapBuilder()
+        .put(1, Arrays.asList(1, 2))
+        .put(2, Arrays.asList(1, 2, 3))
+        .put(3, Arrays.asList()).build()));
+    assertEquals(Arrays.asList(1, 1, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 3), traversal.fun(1).takeWhile(UP_TO(3)).toList());
+  }
+
+  public void testIndefiniteCyclicInterlacedDfs() {
+    Function<Integer, JBIterable<Integer>> traversal = TreeTraversal.INTERLEAVED_DFS.traversal(
+      new Function<Integer, Iterable<Integer>>() {
+        @Override
+        public Iterable<Integer> fun(Integer integer) {
+          JBIterable<Integer> it = JBIterable.generate(1, INCREMENT).takeWhile(UP_TO(integer + 1));
+          // 1: no repeat
+          return it;
+          // 2: repeat indefinitely: all seq
+          //return JBIterable.generate(it, Functions.id()).flatten(Functions.id());
+          // 3: repeat indefinitely: self-cycle
+          //return it.append(JBIterable.generate(integer, Functions.id()));
+        }
+      });
+    JBIterable<Integer> counts = JBIterable.generate(1, INCREMENT).transform(new Function<Integer, Integer>() {
+      @Override
+      public Integer fun(Integer integer) {
+        return traversal.fun(1).takeWhile(UP_TO(integer)).size();
+      }
+    });
+    // 1: no repeat
+    assertEquals(Arrays.asList(1, 4, 13, 39, 117, 359, 1134, 3686, 12276, 41708), counts.take(10).toList());
+    // 2: repeat all seq
+    //assertEquals(Arrays.asList(1, 4, 19, 236), counts.take(4).toList());
+    // 2: repeat self-cycle
+    //assertEquals(Arrays.asList(1, 4, 19, 236), counts.take(4).toList());
+  }
+
   public void testSimplePreOrderDfsBacktrace() {
     List<Integer> backDfs = Collections.emptyList();
     for (TreeTraversal.TracingIt<Integer> it = numTraverser2(TreeTraversal.PRE_ORDER_DFS).fun(1).typedIterator(); it.hasNext(); ) {

platform/util/src/com/intellij/util/containers/TreeTraversal.java

@@ -104,16 +104,28 @@ public abstract class TreeTraversal {
   }
 
   public static abstract class It<T> extends JBIterator<T> {
+    protected final Function<T, ? extends Iterable<? extends T>> tree;
 
+    protected It(Function<T, ? extends Iterable<? extends T>> tree) {
+      this.tree = tree;
+    }
   }
 
   public static abstract class TracingIt<T> extends It<T> {
 
     @Nullable
-    public abstract T parent();
+    public T parent() {
+      throw new UnsupportedOperationException();
+    }
 
     @NotNull
-    public abstract JBIterable<T> backtrace();
+    public JBIterable<T> backtrace() {
+      throw new UnsupportedOperationException();
+    }
+
+    protected TracingIt(Function<T, ? extends Iterable<? extends T>> tree) {
+      super(tree);
+    }
   }
 
   public static abstract class GuidedIt<T> extends It<T> {
@@ -126,9 +138,12 @@ public abstract class TreeTraversal {
     public abstract GuidedIt<T> setGuide(Consumer<GuidedIt<T>> guide);
 
     public abstract GuidedIt<T> queueNext(T child);
-    public abstract GuidedIt<T> queueLast(T child);
     public abstract GuidedIt<T> result(T node);
+    public abstract GuidedIt<T> queueLast(T child);
 
+    protected GuidedIt(Function<T, ? extends Iterable<? extends T>> tree) {
+      super(tree);
+    }
   }
 
   @NotNull
@@ -182,6 +197,24 @@ public abstract class TreeTraversal {
     }
   };
 
+  /**
+   * Returns an iterator over the nodes in a tree structure, using interlaced pre-order
+   * traversal. That is, all paths are traversed in an interlaced manner that is suitable
+   * for infinite and cyclic graphs
+   * and each node's subtrees are traversed before the node itself is returned.
+   * <p/>
+   * <p>No guarantees are made about the behavior of the traversal when nodes change while
+   * iteration is in progress or when the iterators generated by {@code tree} are advanced.
+   */
+  @NotNull
+  public static final TreeTraversal INTERLEAVED_DFS = new TreeTraversal("INTERLEAVED_DFS") {
+    @NotNull
+    @Override
+    public <T> It<T> createIterator(@NotNull Iterable<? extends T> roots, @NotNull Function<T, ? extends Iterable<? extends T>> tree) {
+      return new InterleavedIt<T>(roots, tree);
+    }
+  };
+
   /**
    * Returns an iterator over the nodes in a tree structure, using breadth-first
    * traversal. That is, all the nodes of depth 0 are returned, then depth 1, then 2, and so on.
@@ -216,125 +249,160 @@ public abstract class TreeTraversal {
     }
   };
 
-
   // -----------------------------------------------------------------------------
   // Iterators: DFS
   // -----------------------------------------------------------------------------
 
-  private abstract static class DfsIt<T> extends TracingIt<T> {
-    final ArrayDeque<P<T>> stack = new ArrayDeque<P<T>>();
+  private abstract static class DfsIt<T, H extends P<T, H>> extends TracingIt<T> {
+
+    H last;
+
+    protected DfsIt(Function<T, ? extends Iterable<? extends T>> tree) {
+      super(tree);
+    }
 
     @Nullable
     public T parent() {
-      if (stack.isEmpty()) throw new NoSuchElementException();
-      Iterator<P<T>> it = stack.descendingIterator();
-      it.next();
-      return it.hasNext() ? it.next().node : null;
+      if (last == null) throw new NoSuchElementException();
+
+      H p = last.parent;
+      return p == null ? null : p.node;
     }
 
     @NotNull
     public JBIterable<T> backtrace() {
-      if (stack.isEmpty()) throw new NoSuchElementException();
-      return new JBIterable<P<T>>() {
-        @Override
-        public Iterator<P<T>> iterator() {
-          return stack.descendingIterator();
-        }
-      }.transform(P.<T>toNode()).filter(Condition.NOT_NULL);
+      if (last == null) throw new NoSuchElementException();
+      return JBIterable.generate(last, P.<T>toPrev()).transform(P.<T>toNode()).filter(Condition.NOT_NULL);
     }
   }
 
-  private final static class PreOrderIt<T> extends DfsIt<T> {
-
-    final Function<T, ? extends Iterable<? extends T>> tree;
+  private final static class PreOrderIt<T> extends DfsIt<T, P1<T>> {
 
     PreOrderIt(@NotNull Iterable<? extends T> roots, Function<T, ? extends Iterable<? extends T>> tree) {
-      this.tree = tree;
-      stack.addLast(P.create(roots));
+      super(tree);
+      last = P1.create(roots);
     }
 
     @Override
     public T nextImpl() {
-      while (!stack.isEmpty()) {
-        Iterator<? extends T> it = stack.getLast().iterator(tree);
+      while (last != null) {
+        Iterator<? extends T> it = last.iterator(tree);
         if (it.hasNext()) {
           T result = it.next();
-          stack.addLast(P.create(result));
+          last = last.add(P1.create(result));
           return result;
         }
         else {
-          stack.removeLast();
+          last = last.remove();
         }
       }
       return stop();
     }
   }
 
-  private static final class PostOrderIt<T> extends DfsIt<T> {
-
-    final Function<T, ? extends Iterable<? extends T>> tree;
+  private static final class PostOrderIt<T> extends DfsIt<T, P1<T>> {
 
     PostOrderIt(@NotNull Iterable<? extends T> roots, Function<T, ? extends Iterable<? extends T>> tree) {
-      this.tree = tree;
+      super(tree);
       for (T root : roots) {
-        stack.addLast(P.create(root));
+        P1<T> p = P1.create(root);
+        last = last == null ? p : last.add(p);
       }
     }
 
     @Override
     public T nextImpl() {
-      while (!stack.isEmpty()) {
-        Iterator<? extends T> it = stack.getLast().iterator(tree);
+      while (last != null) {
+        Iterator<? extends T> it = last.iterator(tree);
         if (it.hasNext()) {
           T result = it.next();
-          stack.addLast(P.create(result));
+          last = last.add(P1.create(result));
         }
         else {
-          return stack.removeLast().node;
+          T result = last.node;
+          last = last.remove();
+          return result;
         }
       }
       return stop();
     }
   }
 
-  private final static class LeavesDfsIt<T> extends DfsIt<T> {
-
-    final Function<T, ? extends Iterable<? extends T>> tree;
+  private final static class LeavesDfsIt<T> extends DfsIt<T, P1<T>> {
 
     LeavesDfsIt(@NotNull Iterable<? extends T> roots, Function<T, ? extends Iterable<? extends T>> tree) {
-      this.tree = tree;
-      stack.addLast(P.create(roots));
+      super(tree);
+      last = P1.create(roots);
     }
 
     @Override
     public T nextImpl() {
-      while (!stack.isEmpty()) {
-        P<T> top = stack.getLast();
+      while (last != null) {
+        P1<T> top = last;
         if (top.iterator(tree).hasNext() && !top.empty) {
           T child = top.iterator(tree).next();
-          stack.addLast(P.create(child));
+          last = last.add(P1.create(child));
         }
         else {
-          stack.removeLast();
-          if (top.empty) return stack.isEmpty() ? stop() : top.node;
+          last = last.remove();
+          if (top.empty) return last == null ? stop() : top.node;
         }
       }
       return stop();
     }
   }
 
+  private final static class InterleavedIt<T> extends DfsIt<T, P2<T>> {
+
+    P2<T> cur, max;
+
+    InterleavedIt(@NotNull Iterable<? extends T> roots, Function<T, ? extends Iterable<? extends T>> tree) {
+      super(tree);
+      last = P2.create(roots);
+      cur = max = last;
+    }
+
+    @Override
+    public T nextImpl() {
+      while (last != null) {
+        if (cur == null) {
+          cur = max;
+          max = max.next;
+        }
+        Iterator<? extends T> it = cur.iterator(tree);
+        if (it.hasNext()) {
+          T result = it.next();
+          last = last.add(P2.create(result));
+          last.parent = cur;
+          cur = cur.prev;
+          if (max == null) {
+            max = last;
+          }
+          return result;
+        }
+        else {
+          if (cur == last) {
+            last = cur.prev;
+          }
+          cur = cur.remove();
+        }
+      }
+      return stop();
+    }
+  }
+
+
   // -----------------------------------------------------------------------------
   // Iterators: BFS
   // -----------------------------------------------------------------------------
 
   private static final class PlainBfsIt<T> extends It<T> {
 
-    final Function<T, ? extends Iterable<? extends T>> tree;
     final ArrayDeque<T> queue = new ArrayDeque<T>();
-    P<T> top;
+    P1<T> top;
 
     PlainBfsIt(@NotNull Iterable<? extends T> roots, Function<T, ? extends Iterable<? extends T>> tree) {
-      this.tree = tree;
+      super(tree);
       JBIterable.from(roots).addAllTo(queue);
     }
 
@@ -345,18 +413,17 @@ public abstract class TreeTraversal {
         top = null;
       }
       if (queue.isEmpty()) return stop();
-      top = P.create(queue.remove());
+      top = P1.create(queue.remove());
       return top.node;
     }
   }
 
-  private static final class LeavesBfsIt<T> extends It<T> {
+  private static final class LeavesBfsIt<T> extends TracingIt<T> {
 
-    final Function<T, ? extends Iterable<? extends T>> tree;
     final ArrayDeque<T> queue = new ArrayDeque<T>();
 
     LeavesBfsIt(@NotNull Iterable<? extends T> roots, Function<T, ? extends Iterable<? extends T>> tree) {
-      this.tree = tree;
+      super(tree);
       JBIterable.from(roots).addAllTo(queue);
     }
 
@@ -375,13 +442,12 @@ public abstract class TreeTraversal {
 
   private final static class TracingBfsIt<T> extends TracingIt<T> {
 
-    final Function<T, ? extends Iterable<? extends T>> tree;
     final ArrayDeque<T> queue = new ArrayDeque<T>();
     final Map<T, T> paths = ContainerUtil.newTroveMap(ContainerUtil.<T>identityStrategy());
-    P<T> top;
+    P1<T> top;
 
     TracingBfsIt(@NotNull Iterable<? extends T> roots, Function<T, ? extends Iterable<? extends T>> tree) {
-      this.tree = tree;
+      super(tree);
       JBIterable.from(roots).addAllTo(queue);
     }
 
@@ -396,7 +462,7 @@ public abstract class TreeTraversal {
         top = null;
       }
       if (queue.isEmpty()) return stop();
-      top = P.create(queue.remove());
+      top = P1.create(queue.remove());
       return top.node;
     }
 
@@ -434,15 +500,14 @@ public abstract class TreeTraversal {
   // Misc
   // -----------------------------------------------------------------------------
   private static final class GuidedItImpl<T> extends GuidedIt<T> {
-    final ArrayDeque<P<T>> stack = new ArrayDeque<P<T>>();
-    final Function<T, ? extends Iterable<? extends T>> tree;
+    P1<T> first, last;
 
     Consumer<GuidedIt<T>> guide;
     T curResult;
 
     GuidedItImpl(@NotNull Iterable<? extends T> roots, Function<T, ? extends Iterable<? extends T>> tree) {
-      this.tree = tree;
-      stack.addLast(P.create(roots));
+      super(tree);
+      first = last = P1.create(roots);
     }
 
     public GuidedIt<T> setGuide(Consumer<GuidedIt<T>> guide) {
@@ -451,12 +516,12 @@ public abstract class TreeTraversal {
     }
 
     public GuidedIt<T> queueNext(T child) {
-      if (child != null) stack.addLast(P.create(child));
+      if (child != null) last = last.add(P1.create(child));
       return this;
     }
 
     public GuidedIt<T> queueLast(T child) {
-      if (child != null) stack.addFirst(P.create(child));
+      if (child != null) first = first.addBefore(P1.create(child));
       return this;
     }
 
@@ -468,8 +533,8 @@ public abstract class TreeTraversal {
     @Override
     public T nextImpl() {
       if (guide == null) return stop();
-      while (!stack.isEmpty()) {
-        P<T> top = stack.getLast();
+      while (last != null) {
+        P<T, ?> top = last;
         Iterator<? extends T> it = top.iterator(tree);
         boolean hasNext = it.hasNext();
         curResult = null;
@@ -481,7 +546,7 @@ public abstract class TreeTraversal {
           guide.consume(this);
         }
         if (!hasNext) {
-          stack.removeLast();
+          last = last.remove();
         }
         if (curResult != null) {
           return curResult;
@@ -491,45 +556,113 @@ public abstract class TreeTraversal {
     }
   }
 
-  private static class P<T> {
+  private static class P<T, Self extends P<T, Self>> {
     T node;
     Iterable<? extends T> itle;
     Iterator<? extends T> it;
     boolean empty;
 
-    Iterator<? extends T> iterator(@NotNull Function<T, ? extends Iterable<? extends T>> tree) {
+    Self parent;
+
+    static <T, Self extends P<T, Self>> Self create(Self p, T node) {
+      p.node = node;
+      return p;
+    }
+
+    static <T, Self extends P<T, Self>> Self create(Self p, Iterable<? extends T> it) {
+      p.itle = it;
+      return p;
+    }
+
+
+    final Iterator<? extends T> iterator(@NotNull Function<T, ? extends Iterable<? extends T>> tree) {
       if (it != null) return it;
       it = iterable(tree).iterator();
       empty = itle == null || !it.hasNext();
       return it;
     }
 
-    Iterable<? extends T> iterable(@NotNull Function<T, ? extends Iterable<? extends T>> tree) {
+    final Iterable<? extends T> iterable(@NotNull Function<T, ? extends Iterable<? extends T>> tree) {
       return itle != null ? itle : JBIterable.from(itle = tree.fun(node));
     }
 
-    static <T> P<T> create(T node) {
-      P<T> p = new P<T>();
-      p.node = node;
-      return p;
-    }
+    /** @noinspection unchecked */
+    static <T> Function<P<T, ?>, T> toNode() { return TO_NODE; }
+    /** @noinspection unchecked */
+    static <T> Function<P<T, ?>, P<T, ?>> toPrev() { return TO_PREV; }
 
-    static <T> P<T> create(Iterable<? extends T> it) {
-      P<T> p = new P<T>();
-      p.itle = it;
-      return p;
-    }
-
-    static <T> Function<P<T>, T> toNode() {
-      //noinspection unchecked
-      return TO_NODE;
-    }
-
-    static final Function TO_NODE = new Function<P<?>, Object>() {
+    static final Function TO_NODE = new Function<P<?, ?>, Object>() {
       @Override
-      public Object fun(P<?> tp) {
+      public Object fun(P<?, ?> tp) {
         return tp.node;
       }
     };
+    static final Function TO_PREV = new Function.Mono<P<?, ?>>() {
+      @Override
+      public P<?, ?> fun(P<?, ?> tp) {
+        return tp.parent;
+      }
+    };
+  }
+
+  private static final class P1<T> extends P<T, P1<T>> {
+
+    static <T> P1<T> create(T node) { return create(new P1<T>(), node); }
+    static <T> P1<T> create(Iterable<? extends T> it) { return create(new P1<T>(), it); }
+
+    P1<T> add(@NotNull P1<T> next) {
+      next.parent = this;
+      return next;
+    }
+
+    P1<T> addBefore(@NotNull P1<T> next) {
+      next.parent = null;
+      this.parent = next;
+      return next;
+    }
+
+    P1<T> remove() {
+      P1<T> p = parent;
+      parent = null;
+      return p;
+    }
+
+    @Override
+    public String toString() {
+      int h = 0;
+      for (P1<T> p = parent; p != null; p = p.parent) h++;
+      return h + ": " + node;
+    }
+  }
+
+  private static final class P2<T> extends P<T, P2<T>> {
+    P2<T> next, prev;
+
+    static <T> P2<T> create(T node) { return create(new P2<T>(), node); }
+    static <T> P2<T> create(Iterable<? extends T> it) { return create(new P2<T>(), it); }
+
+    P2<T> add(@NotNull P2<T> next) {
+      next.next = this.next;
+      next.prev = this;
+      this.next = next;
+      return next;
+    }
+
+    P2<T> remove() {
+      P2<T> p = prev;
+      P2<T> n = next;
+      prev = next = null;
+      if (p != null) p.next = n;
+      if (n != null) n.prev = p;
+      return p;
+    }
+
+    @Override
+    public String toString() {
+      int h = 0, t = 0;
+      for (P2<T> p = prev; p != null; p = p.prev) h++;
+      for (P2<T> p = next; p != null; p = p.next) t++;
+      return h + " of " + (h + t + 1) + ": " + node;
+    }
   }
 }