package jkind.slicing;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import jkind.lustre.Equation;
import jkind.lustre.Expr;
import jkind.lustre.IdExpr;
import jkind.lustre.Node;
import jkind.lustre.VarDecl;

public class LustreSlicer {
	public static Node slice(Node node, DependencyMap depMap) {
		return sliceByKeep(node, getPropertyDependencies(node, depMap));
	}

	private static Set<String> getPropertyDependencies(Node node, DependencyMap depMap) {
		Set<String> keep = new HashSet<>();
		for (String prop : node.properties) {
			keep.addAll(depMap.get(prop));
		}
		return keep;
	}

	private static Node sliceByKeep(Node node, Set<String> keep) {
		List<VarDecl> inputs = sliceVarDecls(node.inputs, keep);
		List<VarDecl> outputs = sliceVarDecls(node.outputs, keep);
		List<VarDecl> locals = sliceVarDecls(node.locals, keep);
		List<Equation> equations = sliceEquations(node.equations, keep);
		List<Expr> assertions = sliceAssertions(node.assertions, keep);
		return new Node(node.location, node.id, inputs, outputs, locals, equations,
				node.properties, assertions);
	}

	private static List<VarDecl> sliceVarDecls(List<VarDecl> decls, Set<String> keep) {
		List<VarDecl> sliced = new ArrayList<>();
		for (VarDecl decl : decls) {
			if (keep.contains(decl.id)) {
				sliced.add(decl);
			}
		}
		return sliced;
	}

	private static List<Equation> sliceEquations(List<Equation> equations, Set<String> keep) {
		List<Equation> sliced = new ArrayList<>();
		for (Equation eq : equations) {
			if (containsAny(keep, eq.lhs)) {
				sliced.add(eq);
			}
		}
		return sliced;
	}
	
	private static List<Expr> sliceAssertions(List<Expr> assertions, Set<String> keep) {
		List<Expr> sliced = new ArrayList<>();
		for (Expr assertion : assertions) {
			Set<String> deps = IdExtractorVisitor.getIds(assertion);
			if (deps.size() > 0 && keep.contains(deps.iterator().next())) {
				sliced.add(assertion);
			}
		}
		return sliced;
	}

	private static boolean containsAny(Set<String> keep, List<IdExpr> lhs) {
		for (IdExpr idExpr : lhs) {
			if (keep.contains(idExpr.id)) {
				return true;
			}
		}
		return false;
	}
}

@@ -255,16 +255,60 @@
 VarDecl;
+%0Aimport jkind.lustre.visitors.AstMapVisitor;
 %0A%0Apublic
@@ -327,16 +327,38 @@
 eSlicer 
+extends AstMapVisitor 
 %7B%0A%09publi
@@ -425,26 +425,25 @@
 urn 
-sliceByKeep(node, 
+new LustreSlicer(
 getP
@@ -476,16 +476,28 @@
 depMap))
+.visit(node)
 ;%0A%09%7D%0A%0A%09p
@@ -729,506 +729,120 @@
 ate 
-static Node sliceByKeep(Node node, Set%3CString%3E keep) %7B%0A%09%09List%3CVarDecl%3E inputs = sliceVarDecls(node.inputs, keep);%0A%09%09List%3CVarDecl%3E outputs = sliceVarDecls(node.outputs, keep);%0A%09%09List%3CVarDecl%3E locals = sliceVarDecls(node.locals,
+final Set%3CString%3E keep;%0A%0A%09private LustreSlicer(Set%3CString%3E
  keep)
-;%0A%09%09List%3CEquation%3E equations = sliceEquations(node.equations, keep);%0A%09%09List%3CExpr%3E assertions = sliceAssertions(node.assertions,
+ %7B%0A%09%09this.keep =
  keep
-)
 ;%0A%09
-%09return new Node(node.location, node.id, inputs, outputs, locals, equations,%0A%09%09%09%09node.properties, assertions);%0A%09%7D%0A%0A%09private static
+%7D%0A%0A%09@Override%0A%09protected
  Lis
@@ -852,21 +852,21 @@
 arDecl%3E 
-slice
+visit
 VarDecls
@@ -885,34 +885,16 @@
 l%3E decls
-, Set%3CString%3E keep
 ) %7B%0A%09%09Li
@@ -1050,30 +1050,36 @@
 d;%0A%09%7D%0A%0A%09
-private static
+@Override%0A%09protected
  List%3CEq
@@ -1086,21 +1086,21 @@
 uation%3E 
-slice
+visit
 Equation
@@ -1125,34 +1125,16 @@
 quations
-, Set%3CString%3E keep
 ) %7B%0A%09%09Li
@@ -1297,25 +1297,30 @@
 %0A%09%7D%0A
-%09%0A%09private static
+%0A%09@Override%0A%09protected
  Lis
@@ -1327,21 +1327,21 @@
 t%3CExpr%3E 
-slice
+visit
 Assertio
@@ -1368,26 +1368,8 @@
 ions
-, Set%3CString%3E keep
 ) %7B%0A

Use AstMapVisitor to clean up LustreSlicer