Work on optimizations

4 files changed, 276 insertions, 196 deletions

diff --git a/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTExpression.java b/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTExpression.java
index 385d827..3552926 100644
--- a/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTExpression.java
+++ b/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTExpression.java
@@ -4,10 +4,7 @@ import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Function;
 
-import org.apache.commons.lang3.StringUtils;
-
 import bjc.dicelang.ComplexDice;
-import bjc.dicelang.CompoundDice;
 import bjc.dicelang.IDiceExpression;
 import bjc.dicelang.ast.nodes.DiceASTType;
 import bjc.dicelang.ast.nodes.IDiceASTNode;
@@ -50,10 +47,9 @@ public class DiceASTExpression implements IDiceExpression {
 	 *            The enviroment to evaluate bindings and such against
 	 * @return The operations to use when collapsing the AST
 	 */
-	private static Map<IDiceASTNode, IOperatorCollapser>
-			buildOperations(Map<String, DiceASTExpression> enviroment) {
-		Map<IDiceASTNode, IOperatorCollapser> operatorCollapsers =
-				new HashMap<>();
+	private static Map<IDiceASTNode, IOperatorCollapser> buildOperations(
+			Map<String, DiceASTExpression> enviroment) {
+		Map<IDiceASTNode, IOperatorCollapser> operatorCollapsers = new HashMap<>();
 
 		operatorCollapsers.put(OperatorDiceNode.ADD,
 				(leftNode, rightNode) -> {
@@ -133,9 +129,8 @@ public class DiceASTExpression implements IDiceExpression {
 
 				GenHolder<Boolean> selfReference = new GenHolder<>(false);
 
-				DiceASTReferenceChecker refChecker =
-						new DiceASTReferenceChecker(selfReference,
-								variableName);
+				DiceASTReferenceChecker refChecker = new DiceASTReferenceChecker(
+						selfReference, variableName);
 
 				rightAST.traverse(TreeLinearizationMethod.PREORDER,
 						refChecker);
@@ -168,8 +163,8 @@ public class DiceASTExpression implements IDiceExpression {
 			Pair<Integer, AST<IDiceASTNode>> rightNode) {
 		return leftNode.merge((leftValue, leftAST) -> {
 			return rightNode.merge((rightValue, rightAST) -> {
-				int compoundValue =
-						Integer.parseInt(Integer.toString(leftValue)
+				int compoundValue = Integer
+						.parseInt(Integer.toString(leftValue)
 								+ Integer.toString(rightValue));
 
 				return new Pair<>(compoundValue, new AST<>(
@@ -195,8 +190,8 @@ public class DiceASTExpression implements IDiceExpression {
 									+ rightAST);
 				}
 
-				int rolledValue =
-						new ComplexDice(leftValue, rightValue).roll();
+				int rolledValue = new ComplexDice(leftValue, rightValue)
+						.roll();
 
 				return new Pair<>(rolledValue, new AST<>(
 						OperatorDiceNode.GROUP, leftAST, rightAST));
@@ -231,62 +226,37 @@ public class DiceASTExpression implements IDiceExpression {
 	/**
 	 * Expand a leaf AST token into a pair for evaluation
 	 * 
-	 * @param tokn
+	 * @param leafNode
 	 *            The token to evaluate
 	 * @return A pair consisting of the token's value and the AST it
 	 *         represents
 	 */
-	private Pair<Integer, AST<IDiceASTNode>> evalLeaf(IDiceASTNode tokn) {
-		if (tokn.getType() == DiceASTType.VARIABLE) {
-			return parseVariable(tokn);
-		} else if (tokn.getType() == DiceASTType.LITERAL) {
-			return parseLiteral(tokn);
+	private Pair<Integer, AST<IDiceASTNode>> evaluateLeaf(
+			IDiceASTNode leafNode) {
+		if (leafNode.getType() == DiceASTType.VARIABLE) {
+			VariableDiceNode node = (VariableDiceNode) leafNode;
+
+			return parseVariable(node);
+		} else if (leafNode.getType() == DiceASTType.LITERAL) {
+			LiteralDiceNode node = (LiteralDiceNode) leafNode;
+
+			return node.toParseValue();
 		} else {
 			throw new UnsupportedOperationException("Found leaf operator "
-					+ tokn + ". These aren't supported.");
+					+ leafNode + ". These aren't supported.");
 		}
 	}
 
-	private static Pair<Integer, AST<IDiceASTNode>>
-			parseLiteral(IDiceASTNode tokn) {
-		LiteralDiceNode literalNode = (LiteralDiceNode) tokn;
-		String dat = literalNode.getData();
-
-		if (isValidInfixOperator(dat, "c")) {
-			String[] strangs = dat.split("c");
-
-			return new Pair<>(new CompoundDice(strangs).roll(),
-					new AST<>(tokn));
-		} else if (isValidInfixOperator(dat, "d")) {
-			/*
-			 * Handle dice groups
-			 */
-			return new Pair<>(ComplexDice.fromString(dat).roll(),
-					new AST<>(tokn));
-		} else {
-			try {
-				return new Pair<>(Integer.parseInt(dat), new AST<>(tokn));
-			} catch (NumberFormatException nfex) {
-				throw new UnsupportedOperationException(
-						"Found malformed leaf token " + tokn);
-			}
-		}
-	}
-
-	private static boolean isValidInfixOperator(String dat, String op) {
-		return StringUtils.countMatches(dat, op) == 1
-				&& !dat.equalsIgnoreCase(op) && !dat.startsWith(op);
-	}
-
-	private Pair<Integer, AST<IDiceASTNode>>
-			parseVariable(IDiceASTNode tokn) {
-		String varName = ((VariableDiceNode) tokn).getVariable();
+	private Pair<Integer, AST<IDiceASTNode>> parseVariable(
+			VariableDiceNode leafNode) {
+		String varName = leafNode.getVariable();
 
 		if (env.containsKey(varName)) {
-			return new Pair<>(env.get(varName).roll(), new AST<>(tokn));
+			return new Pair<>(env.get(varName).roll(),
+					new AST<>(leafNode));
 		} else {
 			// Handle special case for defining variables
-			return new Pair<>(0, new AST<>(tokn));
+			return new Pair<>(0, new AST<>(leafNode));
 		}
 	}
 
@@ -306,10 +276,10 @@ public class DiceASTExpression implements IDiceExpression {
 	 */
 	@Override
 	public int roll() {
-		Map<IDiceASTNode, IOperatorCollapser> operations =
-				buildOperations(env);
+		Map<IDiceASTNode, IOperatorCollapser> operations = buildOperations(
+				env);
 
-		return ast.collapse(this::evalLeaf, operations::get,
+		return ast.collapse(this::evaluateLeaf, operations::get,
 				(returnedValue) -> returnedValue
 						.merge((left, right) -> left));
 	}
@@ -323,4 +293,15 @@ public class DiceASTExpression implements IDiceExpression {
 	public String toString() {
 		return ast.toString();
 	}
+
+	@Override
+	public int optimize() {
+		throw new UnsupportedOperationException(
+				"Use DiceASTOptimizer for optimizing these");
+	}
+
+	@Override
+	public boolean canOptimize() {
+		return false;
+	}
 }
\ No newline at end of file
diff --git a/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTFreezer.java b/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTFreezer.java
deleted file mode 100644
index 6fa4883..0000000
--- a/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTFreezer.java
+++ /dev/null
@@ -1,136 +0,0 @@
-package bjc.dicelang.ast;
-
-import java.util.function.Function;
-
-import bjc.dicelang.ast.nodes.DiceASTType;
-import bjc.dicelang.ast.nodes.IDiceASTNode;
-import bjc.dicelang.ast.nodes.VariableDiceNode;
-import bjc.utils.funcdata.FunctionalList;
-import bjc.utils.funcdata.FunctionalMap;
-import bjc.utils.parserutils.AST;
-
-/**
- * Freeze references in a dice AST, replacing variable references with what
- * the variables refer to
- * 
- * @author ben
- *
- */
-public class DiceASTFreezer {
-	private static class NodeFreezer
-			implements Function<IDiceASTNode, AST<IDiceASTNode>> {
-		private FunctionalMap<String, AST<IDiceASTNode>> enviroment;
-
-		public NodeFreezer(FunctionalMap<String, AST<IDiceASTNode>> env) {
-			enviroment = env;
-		}
-
-		@Override
-		public AST<IDiceASTNode> apply(IDiceASTNode nod) {
-			if (nod.getType() == DiceASTType.VARIABLE) {
-				return expandNode((VariableDiceNode) nod);
-			} else {
-				return new AST<>(nod);
-			}
-		}
-
-		protected AST<IDiceASTNode>
-				expandNode(VariableDiceNode variableNode) {
-			String varName = variableNode.getVariable();
-
-			if (!enviroment.containsKey(varName)) {
-				throw new IllegalArgumentException(
-						"Attempted to freeze reference"
-								+ " to an undefined variable " + varName);
-			}
-
-			return enviroment.get(varName);
-		}
-	}
-
-	private static final class SelectiveFreezer extends NodeFreezer {
-
-		private FunctionalList<String> variableNames;
-
-		public SelectiveFreezer(
-				FunctionalMap<String, AST<IDiceASTNode>> env,
-				FunctionalList<String> varNames) {
-			super(env);
-			
-			variableNames = varNames;
-		}
-
-		@Override
-		protected AST<IDiceASTNode>
-				expandNode(VariableDiceNode variableNode) {
-			if (variableNames.contains(variableNode.getVariable())) {
-				return super.expandNode(variableNode);
-			} else {
-				return new AST<>(variableNode);
-			}
-		}
-	}
-
-	/**
-	 * Freeze the references in an AST
-	 * 
-	 * @param tree
-	 *            The tree to freeze references in
-	 * @param env
-	 *            The enviroment to get reference values from
-	 * @return The tree with references frozen
-	 */
-	public static AST<IDiceASTNode> freezeAST(AST<IDiceASTNode> tree,
-			FunctionalMap<String, AST<IDiceASTNode>> env) {
-		return selectiveFreeze(tree, env);
-	}
-
-	/**
-	 * Freeze the references in an expression backed by an AST
-	 * 
-	 * @param tree
-	 *            The tree-backed expression to freeze references in
-	 * @param env
-	 *            The enviroment to get reference values from
-	 * @return The tree with references frozen
-	 */
-	public static AST<IDiceASTNode> freezeAST(DiceASTExpression tree,
-			FunctionalMap<String, DiceASTExpression> env) {
-		return freezeAST(tree.getAst(),
-				env.mapValues(expression -> expression.getAst()));
-	}
-
-	/**
-	 * Freeze references to specified variables
-	 * 
-	 * @param tree
-	 *            The tree-backed expression to freeze references in
-	 * @param env
-	 *            The enviroment to resolve variables against
-	 * @param varNames
-	 *            The names of the variables to freeze
-	 * @return An AST with the specified variables frozen
-	 */
-	public static AST<IDiceASTNode> selectiveFreeze(AST<IDiceASTNode> tree,
-			FunctionalMap<String, AST<IDiceASTNode>> env,
-			String... varNames) {
-		return selectiveFreeze(tree, env, new FunctionalList<>(varNames));
-	}
-
-	/**
-	 * Freeze references to specified variables
-	 * 
-	 * @param tree
-	 *            The tree-backed expression to freeze references in
-	 * @param env
-	 *            The enviroment to resolve variables against
-	 * @param varNames
-	 *            The names of the variables to freeze
-	 * @return An AST with the specified variables frozen
-	 */
-	public static AST<IDiceASTNode> selectiveFreeze(AST<IDiceASTNode> tree,
-			FunctionalMap<String, AST<IDiceASTNode>> env,
-			FunctionalList<String> varNames) {
-		return tree.expand(new SelectiveFreezer(env, varNames));
-	}
-}
\ No newline at end of file
diff --git a/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTOptimizer.java b/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTOptimizer.java
new file mode 100644
index 0000000..f5a6696
--- /dev/null
+++ b/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTOptimizer.java
@@ -0,0 +1,144 @@
+package bjc.dicelang.ast;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.BiFunction;
+import java.util.function.BinaryOperator;
+
+import bjc.dicelang.IDiceExpression;
+import bjc.dicelang.ast.nodes.DiceASTType;
+import bjc.dicelang.ast.nodes.IDiceASTNode;
+import bjc.dicelang.ast.nodes.LiteralDiceNode;
+import bjc.dicelang.ast.nodes.OperatorDiceNode;
+
+import static bjc.dicelang.ast.nodes.DiceASTType.*;
+
+import bjc.utils.parserutils.AST;
+
+/**
+ * Optimize an AST
+ * 
+ * @author ben
+ *
+ */
+public class DiceASTOptimizer {
+	private static final class ArithmeticOperationCollapser
+			implements BinaryOperator<AST<IDiceASTNode>> {
+		private IDiceASTNode							type;
+		private BiFunction<Integer, Integer, Integer>	valueCollapser;
+		private boolean									doSwap;
+
+		public ArithmeticOperationCollapser(IDiceASTNode type,
+				BiFunction<Integer, Integer, Integer> valueCollapser,
+				boolean doSwap) {
+			this.type = type;
+			this.valueCollapser = valueCollapser;
+			this.doSwap = doSwap;
+		}
+
+		@Override
+		public AST<IDiceASTNode> apply(AST<IDiceASTNode> leftAST,
+				AST<IDiceASTNode> rightAST) {
+			boolean isLeftConstant =
+					DiceASTOptimizer.checkNodeType(leftAST, LITERAL)
+							&& DiceASTOptimizer.isNodeConstant(leftAST);
+
+			boolean isRightConstant =
+					DiceASTOptimizer.checkNodeType(rightAST, LITERAL)
+							&& DiceASTOptimizer.isNodeConstant(leftAST);
+
+			if (isLeftConstant) {
+				if (isRightConstant) {
+					int combinedValue = valueCollapser.apply(
+							getNodeValue(leftAST), getNodeValue(rightAST));
+
+					return new AST<>(new LiteralDiceNode(combinedValue));
+				}
+
+				if (doSwap) {
+					return new AST<>(type, rightAST, leftAST);
+				}
+			}
+
+			return new AST<>(type, leftAST, rightAST);
+		}
+	}
+
+	private static Map<IDiceASTNode, BinaryOperator<AST<IDiceASTNode>>>
+			buildConstantCollapsers() {
+		Map<IDiceASTNode, BinaryOperator<AST<IDiceASTNode>>> operatorCollapsers =
+				new HashMap<>();
+
+		operatorCollapsers.put(OperatorDiceNode.ADD,
+				new ArithmeticOperationCollapser(OperatorDiceNode.ADD,
+						(leftVal, rightVal) -> leftVal + rightVal, true));
+		operatorCollapsers.put(OperatorDiceNode.SUBTRACT,
+				new ArithmeticOperationCollapser(OperatorDiceNode.SUBTRACT,
+						(leftVal, rightVal) -> leftVal - rightVal, false));
+		operatorCollapsers.put(OperatorDiceNode.DIVIDE,
+				new ArithmeticOperationCollapser(OperatorDiceNode.DIVIDE,
+						(leftVal, rightVal) -> leftVal / rightVal, false));
+		operatorCollapsers.put(OperatorDiceNode.MULTIPLY,
+				new ArithmeticOperationCollapser(OperatorDiceNode.MULTIPLY,
+						(leftVal, rightVal) -> leftVal * rightVal, true));
+
+		return null;
+	}
+
+	private static AST<IDiceASTNode> collapseLeaf(IDiceASTNode leaf) {
+		// Can't optimize a simple reference
+		if (leaf.getType() == VARIABLE) {
+			return new AST<>(leaf);
+		} else if (leaf.getType() == LITERAL) {
+			LiteralDiceNode node = (LiteralDiceNode) leaf;
+
+			return new AST<>(optimizeLiteral(node, node.toExpression()));
+		} else {
+			throw new UnsupportedOperationException(
+					"Found leaf operator. This isn't supported");
+		}
+	}
+
+	private static IDiceASTNode optimizeLiteral(LiteralDiceNode node,
+			IDiceExpression leaf) {
+		if (leaf.canOptimize()) {
+			return new LiteralDiceNode(Integer.toString(leaf.optimize()));
+		} else {
+			return node;
+		}
+	}
+
+	private static AST<IDiceASTNode> finishTree(AST<IDiceASTNode> tree) {
+		return tree;
+	}
+
+	/**
+	 * Optimize a tree of expressions into a simpler form
+	 * 
+	 * @param tree
+	 *            The tree to optimize
+	 * @return The optimized tree
+	 */
+	public static AST<IDiceASTNode> optimizeTree(AST<IDiceASTNode> tree) {
+		AST<IDiceASTNode> astWithFoldedConstants =
+				tree.collapse(DiceASTOptimizer::collapseLeaf,
+						buildConstantCollapsers()::get,
+						DiceASTOptimizer::finishTree);
+		return astWithFoldedConstants;
+	}
+
+	private static boolean checkNodeType(AST<IDiceASTNode> ast,
+			DiceASTType type) {
+		return ast.applyToHead((node) -> node.getType()) == type;
+	}
+
+	private static boolean isNodeConstant(AST<IDiceASTNode> ast) {
+		return ast.applyToHead(
+				(node) -> ((LiteralDiceNode) node).isConstant());
+	}
+
+	private static int getNodeValue(AST<IDiceASTNode> ast) {
+		return ast.applyToHead(
+				(node) -> ((LiteralDiceNode) node).toConstant());
+	}
+}
diff --git a/dice-lang/src/main/java/bjc/dicelang/ast/nodes/LiteralDiceNode.java b/dice-lang/src/main/java/bjc/dicelang/ast/nodes/LiteralDiceNode.java
index fe4c402..e689c7f 100644
--- a/dice-lang/src/main/java/bjc/dicelang/ast/nodes/LiteralDiceNode.java
+++ b/dice-lang/src/main/java/bjc/dicelang/ast/nodes/LiteralDiceNode.java
@@ -1,5 +1,14 @@
 package bjc.dicelang.ast.nodes;
 
+import org.apache.commons.lang3.StringUtils;
+
+import bjc.dicelang.ComplexDice;
+import bjc.dicelang.CompoundDice;
+import bjc.dicelang.IDiceExpression;
+import bjc.dicelang.ScalarDie;
+import bjc.utils.data.Pair;
+import bjc.utils.parserutils.AST;
+
 /**
  * A AST node that represents a literal value
  * 
@@ -7,6 +16,11 @@ package bjc.dicelang.ast.nodes;
  *
  */
 public class LiteralDiceNode implements IDiceASTNode {
+	private static boolean isValidInfixOperator(String dat, String op) {
+		return StringUtils.countMatches(dat, op) == 1
+				&& !dat.equalsIgnoreCase(op) && !dat.startsWith(op);
+	}
+
 	/**
 	 * The value contained by this node
 	 */
@@ -22,6 +36,14 @@ public class LiteralDiceNode implements IDiceASTNode {
 		this.value = data;
 	}
 
+	/**
+	 * Create a new node with the given value
+	 * @param val The value for this node
+	 */
+	public LiteralDiceNode(int val) {
+		this(Integer.toString(val));
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * 
@@ -82,6 +104,48 @@ public class LiteralDiceNode implements IDiceASTNode {
 		return false;
 	}
 
+	/**
+	 * Parse this node into an expression
+	 * 
+	 * @return The node in expression form
+	 */
+	public IDiceExpression toExpression() {
+		String literalData = this.getData();
+
+		if (LiteralDiceNode.isValidInfixOperator(literalData, "c")) {
+			String[] strangs = literalData.split("c");
+
+			return new CompoundDice(strangs);
+		} else if (LiteralDiceNode.isValidInfixOperator(literalData,
+				"d")) {
+			/*
+			 * Handle dice groups
+			 */
+			return ComplexDice.fromString(literalData);
+		} else {
+			try {
+				return new ScalarDie(Integer.parseInt(literalData));
+			} catch (NumberFormatException nfex) {
+				throw new UnsupportedOperationException(
+						"Found malformed leaf token " + this);
+			}
+		}
+	}
+
+	/**
+	 * Parse this node into an expression
+	 * 
+	 * @return The node as a pair of a sample value and the AST it
+	 *         represents
+	 */
+	public Pair<Integer, AST<IDiceASTNode>> toParseValue() {
+		AST<IDiceASTNode> returnedAST = new AST<>(this);
+
+		IDiceExpression expression = toExpression();
+
+		return new Pair<>(expression.roll(), returnedAST);
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * 
@@ -91,4 +155,31 @@ public class LiteralDiceNode implements IDiceASTNode {
 	public String toString() {
 		return value;
 	}
+
+	/**
+	 * Check if this node represents a constant value
+	 * 
+	 * @return Whether or not this node represents a constant value
+	 */
+	public boolean isConstant() {
+		try {
+			Integer.parseInt(value);
+			return true;
+		} catch (NumberFormatException nfex) {
+			return false;
+		}
+	}
+
+	/**
+	 * Return the constant value this node represents
+	 * 
+	 * @return The constant value of this node
+	 * 
+	 * @throws NumberFormatException
+	 *             if you call this on a node that doesn't represent a
+	 *             constant value
+	 */
+	public int toConstant() {
+		return Integer.parseInt(value);
+	}
 }
\ No newline at end of file