ExpressionTraverseNonNested: handle strategy, CTL operators as well

We don't want to recurse into the strategy or CTL operators.

Refactor a bit to keep the boilerplate code for each visit() method down.

prism/src/parser/visitor/ExpressionTraverseNonNested.java

@@ -26,26 +26,33 @@
 
 package parser.visitor;
 
+import parser.ast.ExpressionExists;
+import parser.ast.ExpressionForAll;
 import parser.ast.ExpressionProb;
 import parser.ast.ExpressionReward;
 import parser.ast.ExpressionSS;
+import parser.ast.ExpressionStrategy;
 import prism.PrismLangException;
 
 /**
- * Base class for recursively traversing an Expression, but not
- * entering nested P/R/SS subexpressions.
- *
- * By default, will not recurse into the P/R/SS elements, but
- * can be configured to allow a certain amount of nested P/R/SS
- * elements before stopping.
- *
- * <br/>
- * Subclasses should not override<br/>
+ * Base class for recursively traversing an Expression,
+ * but not recursing into nested state formulas, i.e.,
  * <ul>
- *  <li>{@code visit(ExpressionProb)}</li>
- *  <li>{@code visit(ExpressionReward)}</li>
- *  <li>{@code visit(ExpressionSS)}</li>
+ * <li>the P, R, S operators,</li>
+ * <li>the CTL E and A operators,</li>
+ * <li>the strategy &lt;&lt;>> and [[]] operators.</li>
  * </ul>
+ * <br>
+ * Intuitively, given a PRCTL* path formula, all expressions
+ * that form the top-level LTL path formula will be visited,
+ * without recursing into the nested non-Boolean state formulas.
+ * <br>
+ * Optionally, a nesting limit can be specified to allow a given
+ * amount of nesting levels.
+ *
+ * <br/>
+ * Subclasses should not override any of the {@code visit(...)} methods;
+ * override the {@code visitPre()} or {@code visitPost()} methods instead.
  */
 public class ExpressionTraverseNonNested extends ASTTraverse
 {
@@ -73,45 +80,80 @@ public class ExpressionTraverseNonNested extends ASTTraverse
 		return currentNesting <= nestingLimit;
 	}
 
-	@Override
-	public Object visit(ExpressionProb e) throws PrismLangException
+	/**
+	 * Check whether entering a nested expression at the current
+	 * nesting level is fine.
+	 */
+	private boolean enter()
 	{
 		currentNesting++;
 		// only visit if we are still in limit
 		if (!inLimit()) {
 			currentNesting--;
-			return null;
+			return false;
 		}
-		Object rv = super.visit(e);
+		return true;
+	}
+
+	/** Exits the current nesting level and return the object resulting from the visit */
+	private Object exit(Object e)
+	{
 		currentNesting--;
-		return rv;
+		return e;
 	}
 
+	// P operator
+	@Override
+	public Object visit(ExpressionProb e) throws PrismLangException
+	{
+		if (!enter())
+			return null;
+		return exit(super.visit(e));
+	}
+
+	// R operator
 	@Override
 	public Object visit(ExpressionReward e) throws PrismLangException
 	{
-		currentNesting++;
-		// only visit if we are still in limit
-		if (!inLimit()) {
-			currentNesting--;
+		if (!enter())
 			return null;
-		}
-		Object rv = super.visit(e);
-		currentNesting--;
-		return rv;
+		return exit(super.visit(e));
 	}
 
+	// S operator
 	@Override
 	public Object visit(ExpressionSS e) throws PrismLangException
 	{
-		currentNesting++;
-		// only visit if we are still in limit
-		if (!inLimit()) {
-			currentNesting--;
+		if (!enter())
 			return null;
-		}
-		Object rv = super.visit(e);
-		currentNesting--;
-		return rv;
+		return exit(super.visit(e));
+	}
+
+	// strategy operators
+	@Override
+	public Object visit(ExpressionStrategy e) throws PrismLangException
+	{
+		if (!enter())
+			return null;
+		return exit(super.visit(e));
 	}
+
+	// CTL E operator
+	@Override
+	public Object visit(ExpressionExists e) throws PrismLangException
+	{
+		if (!enter())
+			return null;
+		return exit(super.visit(e));
+	}
+
+	// CTL A operator
+	@Override
+	public Object visit(ExpressionForAll e) throws PrismLangException
+	{
+		if (!enter())
+			return null;
+		return exit(super.visit(e));
+	}
+
 }