Merge pull request #493 from viperproject/optimize-iterator

Fix performance regression of AST traversal

Author: fpoli

src/main/scala/viper/silver/ast/Ast.scala

@@ -13,7 +13,6 @@ import viper.silver.ast.utility.rewriter.Traverse.Traverse
 import viper.silver.ast.utility.rewriter.{Rewritable, StrategyBuilder, Traverse}
 import viper.silver.verifier.errors.ErrorNode
 import viper.silver.verifier.{AbstractVerificationError, ConsistencyError, ErrorReason}
-import scala.language.reflectiveCalls
 
 /*
 
@@ -49,7 +48,7 @@ Some design choices:
   * Note that all but Program are transitive subtypes of `Node` via `Hashable`. The reason is
   * that AST node hashes may depend on the entire program, not just their sub-AST.
   */
-trait Node extends Iterable[Node] with Rewritable {
+trait Node extends Traversable[Node] with Rewritable {
 
   /** @see [[Nodes.subnodes()]] */
   def subnodes = Nodes.subnodes(this)
@@ -62,23 +61,7 @@ trait Node extends Iterable[Node] with Rewritable {
     Visitor.reduceWithContext(this, Nodes.subnodes)(context, enter, combine)
   }
 
-  override def iterator: Iterator[Node] = {
-    val iterator = new Iterator[Node] {
-      var remaining = Seq.empty[Node]
-
-      override def hasNext: Boolean = remaining.nonEmpty
-
-      override def next(): Node = {
-        val result = remaining.head
-        remaining = remaining.tail
-        result
-      }
-    }
-
-    Visitor.visit(this, Nodes.subnodes) { case n: Node => iterator.remaining ++= Seq(n) }
-
-    iterator
-  }
+  override def foreach[A](f: Node => A) = Visitor.visit(this, Nodes.subnodes) { case a: Node => f(a) }
 
   /** @see [[Visitor.visit()]] */
   def visit[A](f: PartialFunction[Node, A]): Unit = {

src/main/scala/viper/silver/ast/utility/Traversable.scala

@@ -0,0 +1,59 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+//
+// Copyright (c) 2011-2020 ETH Zurich.
+
+package viper.silver.ast.utility
+
+import scala.collection.mutable
+
+/** A trait for traversable collections. */
+trait Traversable[+A] {
+  self =>
+
+  /** Applies a function to all element of the collection. */
+  def foreach[B](f: A => B): Unit
+
+  /** Creates a filter of this traversable collection. */
+  def withFilter(p: A => Boolean): Traversable[A] = new WithFilter(p)
+
+  class WithFilter(p: A => Boolean) extends Traversable[A] {
+    /** Applies a function to all filtered elements of the outer collection. */
+    def foreach[U](f: A => U): Unit = {
+      for (x <- self) {
+        if (p(x)) f(x)
+      }
+    }
+
+    /** Further refines the filter of this collection. */
+    override def withFilter(q: A => Boolean): WithFilter = {
+      new WithFilter(x => p(x) && q(x))
+    }
+  }
+
+  /** Finds the first element of this collection for which the given partial
+    * function is defined, and applies the partial function to it.
+    */
+  def collectFirst[B](pf: PartialFunction[A, B]): Option[B] = {
+    for (x <- self) {
+      if (pf.isDefinedAt(x)) {
+        return Some(pf(x))
+      }
+    }
+    None
+  }
+
+  /** Builds a new collection by applying a partial function to all elements
+    * of this collection on which the function is defined.
+    */
+  def collect[B](pf: PartialFunction[A, B]): Iterable[B] = {
+    val elements = mutable.Queue.empty[B]
+    for (x <- self) {
+      if (pf.isDefinedAt(x)) {
+        elements.append(pf(x))
+      }
+    }
+    elements
+  }
+}