1 files changed, 64 insertions, 37 deletions

diff --git a/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTFlattener.java b/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTFlattener.java
index aa15c4d..fad4870 100644
--- a/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTFlattener.java
+++ b/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTFlattener.java
@@ -3,17 +3,19 @@ package bjc.dicelang.ast;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.BinaryOperator;
+import java.util.function.Function;
 
 import org.apache.commons.lang3.StringUtils;
 
 import bjc.dicelang.BindingDiceExpression;
 import bjc.dicelang.ComplexDice;
 import bjc.dicelang.CompoundDice;
-import bjc.dicelang.CompoundDiceExpression;
+import bjc.dicelang.OperatorDiceExpression;
 import bjc.dicelang.DiceExpressionType;
 import bjc.dicelang.IDiceExpression;
 import bjc.dicelang.ReferenceDiceExpression;
 import bjc.dicelang.ScalarDie;
+
 import bjc.utils.parserutils.AST;
 
 /**
@@ -24,6 +26,60 @@ import bjc.utils.parserutils.AST;
  *
  */
 public class DiceASTFlattener {
+	private static final class NodeCollapser
+			implements Function<IDiceASTNode, IDiceExpression> {
+		private Map<String, IDiceExpression> enviroment;
+
+		public NodeCollapser(Map<String, IDiceExpression> env) {
+			this.enviroment = env;
+		}
+
+		@Override
+		public IDiceExpression apply(IDiceASTNode nod) {
+			if (nod.getType() == DiceASTType.LITERAL) {
+				return expFromLiteral((LiteralDiceNode) nod);
+			} else if (nod.getType() == DiceASTType.VARIABLE) {
+				String varName = ((VariableDiceNode) nod).getVariable();
+
+				return new ReferenceDiceExpression(varName, enviroment);
+			} else {
+				throw new UnsupportedOperationException(
+						"Attempted to flatten something that can't be"
+								+ " flattened. The culprit is " + nod);
+			}
+		}
+
+		/**
+		 * Create a dice expression from a literal token
+		 * 
+		 * @param tok
+		 *            The token to convert to an expression
+		 * @return The dice expression represented by the token
+		 */
+		private static IDiceExpression
+				expFromLiteral(LiteralDiceNode tok) {
+			String data = tok.getData();
+
+			if (data.equals("")) {
+				throw new UnsupportedOperationException(
+						"Can't convert a blank token into a literal");
+			}
+
+			if (StringUtils.countMatches(data, 'c') == 1
+					&& !data.equalsIgnoreCase("c")) {
+				String[] strangs = data.split("c");
+
+				return new CompoundDice(ComplexDice.fromString(strangs[0]),
+						ComplexDice.fromString(strangs[1]));
+			} else if (StringUtils.countMatches(data, 'd') == 1
+					&& !data.equalsIgnoreCase("d")) {
+				return ComplexDice.fromString(data);
+			} else {
+				return new ScalarDie(Integer.parseInt(data));
+			}
+		}
+	}
+
 	/**
 	 * Build the operations to use for tree flattening
 	 * 
@@ -35,20 +91,21 @@ public class DiceASTFlattener {
 			buildOperations(Map<String, IDiceExpression> env) {
 		Map<IDiceASTNode, BinaryOperator<IDiceExpression>> opCollapsers =
 				new HashMap<>();
+
 		opCollapsers.put(OperatorDiceNode.ADD, (left, right) -> {
-			return new CompoundDiceExpression(right, left,
+			return new OperatorDiceExpression(right, left,
 					DiceExpressionType.ADD);
 		});
 		opCollapsers.put(OperatorDiceNode.SUBTRACT, (left, right) -> {
-			return new CompoundDiceExpression(right, left,
+			return new OperatorDiceExpression(right, left,
 					DiceExpressionType.SUBTRACT);
 		});
 		opCollapsers.put(OperatorDiceNode.MULTIPLY, (left, right) -> {
-			return new CompoundDiceExpression(right, left,
+			return new OperatorDiceExpression(right, left,
 					DiceExpressionType.MULTIPLY);
 		});
 		opCollapsers.put(OperatorDiceNode.DIVIDE, (left, right) -> {
-			return new CompoundDiceExpression(right, left,
+			return new OperatorDiceExpression(right, left,
 					DiceExpressionType.DIVIDE);
 		});
 		opCollapsers.put(OperatorDiceNode.ASSIGN, (left, right) -> {
@@ -65,30 +122,6 @@ public class DiceASTFlattener {
 	}
 
 	/**
-	 * Create a dice expression from a literal token
-	 * 
-	 * @param tok
-	 *            The token to convert to an expression
-	 * @return The dice expression represented by the token
-	 */
-	private static IDiceExpression expFromLiteral(LiteralDiceNode tok) {
-		String data = tok.getData();
-
-		if (StringUtils.countMatches(data, 'c') == 1
-				&& !data.equalsIgnoreCase("c")) {
-			String[] strangs = data.split("c");
-
-			return new CompoundDice(ComplexDice.fromString(strangs[0]),
-					ComplexDice.fromString(strangs[1]));
-		} else if (StringUtils.countMatches(data, 'd') == 1
-				&& !data.equalsIgnoreCase("d")) {
-			return ComplexDice.fromString(data);
-		} else {
-			return new ScalarDie(Integer.parseInt(data));
-		}
-	}
-
-	/**
 	 * Flatten a AST into a dice expression
 	 * 
 	 * @param ast
@@ -102,13 +135,7 @@ public class DiceASTFlattener {
 		Map<IDiceASTNode, BinaryOperator<IDiceExpression>> opCollapsers =
 				buildOperations(env);
 
-		return ast.collapse((nod) -> {
-			if (nod instanceof LiteralDiceNode) {
-				return expFromLiteral((LiteralDiceNode) nod);
-			} else {
-				return new ReferenceDiceExpression(
-						((VariableDiceNode) nod).getVariable(), env);
-			}
-		} , opCollapsers::get, (r) -> r);
+		return ast.collapse(new NodeCollapser(env), opCollapsers::get,
+				(r) -> r);
 	}
 }