The optimizer is working :)

2 files changed, 83 insertions, 6 deletions

diff --git a/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTOptimizer.java b/dice-lang/src/main/java/bjc/dicelang/ast/optimization/DiceASTOptimizer.java
index f5a6696..e6c62ee 100644
--- a/dice-lang/src/main/java/bjc/dicelang/ast/DiceASTOptimizer.java
+++ b/dice-lang/src/main/java/bjc/dicelang/ast/optimization/DiceASTOptimizer.java
@@ -1,4 +1,4 @@
-package bjc.dicelang.ast;
+package bjc.dicelang.ast.optimization;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -22,6 +22,49 @@ import bjc.utils.parserutils.AST;
  *
  */
 public class DiceASTOptimizer {
+	private static final class NestedArithmeticOperationCollapser
+			implements BinaryOperator<AST<IDiceASTNode>> {
+		private IDiceASTNode							type;
+		private BiFunction<Integer, Integer, Integer>	valueCollapser;
+
+		public NestedArithmeticOperationCollapser(IDiceASTNode type,
+				BiFunction<Integer, Integer, Integer> valueCollapser) {
+			this.type = type;
+			this.valueCollapser = valueCollapser;
+		}
+
+		@Override
+		public AST<IDiceASTNode> apply(AST<IDiceASTNode> leftAST,
+				AST<IDiceASTNode> rightAST) {
+			AST<IDiceASTNode> rightBranchOfLeftAST =
+					leftAST.applyToRight((rightSideAST) -> rightSideAST);
+			AST<IDiceASTNode> leftBranchOfLeftAST =
+					leftAST.applyToRight((rightSideAST) -> rightSideAST);
+
+			boolean leftContainsNestedConstant = DiceASTOptimizer
+					.checkNodeType(rightBranchOfLeftAST, LITERAL)
+					&& DiceASTOptimizer.isNodeConstant(leftAST);
+
+			boolean isRightConstant =
+					DiceASTOptimizer.checkNodeType(rightAST, LITERAL)
+							&& DiceASTOptimizer.isNodeConstant(leftAST);
+
+			if (leftContainsNestedConstant && isRightConstant) {
+				int combinedValue = valueCollapser.apply(
+						getNodeValue(rightBranchOfLeftAST),
+						getNodeValue(rightAST));
+
+				AST<IDiceASTNode> newRightBranch =
+						new AST<>(new LiteralDiceNode(combinedValue));
+
+				return new AST<>(type, leftBranchOfLeftAST,
+						newRightBranch);
+			}
+
+			return new AST<>(type, leftAST, rightAST);
+		}
+	}
+
 	private static final class ArithmeticOperationCollapser
 			implements BinaryOperator<AST<IDiceASTNode>> {
 		private IDiceASTNode							type;
@@ -72,17 +115,38 @@ public class DiceASTOptimizer {
 		operatorCollapsers.put(OperatorDiceNode.ADD,
 				new ArithmeticOperationCollapser(OperatorDiceNode.ADD,
 						(leftVal, rightVal) -> leftVal + rightVal, true));
+
+		operatorCollapsers.put(OperatorDiceNode.MULTIPLY,
+				new ArithmeticOperationCollapser(OperatorDiceNode.MULTIPLY,
+						(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));
+
+		return operatorCollapsers;
+	}
+
+	private static Map<IDiceASTNode, BinaryOperator<AST<IDiceASTNode>>>
+			buildNestedConstantCollapsers() {
+		Map<IDiceASTNode, BinaryOperator<AST<IDiceASTNode>>> operatorCollapsers =
+				new HashMap<>();
+
+		operatorCollapsers.put(OperatorDiceNode.ADD,
+				new NestedArithmeticOperationCollapser(
+						OperatorDiceNode.ADD,
+						(leftVal, rightVal) -> leftVal + rightVal));
+
 		operatorCollapsers.put(OperatorDiceNode.MULTIPLY,
-				new ArithmeticOperationCollapser(OperatorDiceNode.MULTIPLY,
-						(leftVal, rightVal) -> leftVal * rightVal, true));
+				new NestedArithmeticOperationCollapser(
+						OperatorDiceNode.MULTIPLY,
+						(leftVal, rightVal) -> leftVal * rightVal));
 
-		return null;
+		return operatorCollapsers;
 	}
 
 	private static AST<IDiceASTNode> collapseLeaf(IDiceASTNode leaf) {
@@ -120,11 +184,18 @@ public class DiceASTOptimizer {
 	 * @return The optimized tree
 	 */
 	public static AST<IDiceASTNode> optimizeTree(AST<IDiceASTNode> tree) {
-		AST<IDiceASTNode> astWithFoldedConstants =
+		AST<IDiceASTNode> astWithConstantsFolded =
 				tree.collapse(DiceASTOptimizer::collapseLeaf,
 						buildConstantCollapsers()::get,
 						DiceASTOptimizer::finishTree);
-		return astWithFoldedConstants;
+
+		AST<IDiceASTNode> astWithNestedConstantsFolded =
+				astWithConstantsFolded.collapse(
+						DiceASTOptimizer::collapseLeaf,
+						buildNestedConstantCollapsers()::get,
+						DiceASTOptimizer::finishTree);
+
+		return astWithNestedConstantsFolded;
 	}
 
 	private static boolean checkNodeType(AST<IDiceASTNode> ast,
diff --git a/dice-lang/src/main/java/bjc/dicelang/ast/optimization/package-info.java b/dice-lang/src/main/java/bjc/dicelang/ast/optimization/package-info.java
new file mode 100644
index 0000000..132563d
--- /dev/null
+++ b/dice-lang/src/main/java/bjc/dicelang/ast/optimization/package-info.java
@@ -0,0 +1,6 @@
+/**
+ * Contains classes for optimizing ASTs
+ * @author ben
+ *
+ */
+package bjc.dicelang.ast.optimization;
\ No newline at end of file