Implement precedence in Rust writer

3 files changed, 120 insertions, 17 deletions

diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustAssociativity.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustAssociativity.java
new file mode 100644
index 0000000..632e967
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustAssociativity.java
@@ -0,0 +1,7 @@
+package org.zwobble.hobgoblin.compiler.output.lang.rust;
+
+public enum RustAssociativity {
+    LEFT,
+    RIGHT,
+    NONE;
+}
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustPrecedence.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustPrecedence.java
new file mode 100644
index 0000000..c3038ca
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustPrecedence.java
@@ -0,0 +1,27 @@
+package org.zwobble.hobgoblin.compiler.output.lang.rust;
+
+public enum RustPrecedence {
+    PRIMARY,
+    PATHS,
+    METHOD_CALLS,
+    FIELD_EXPRESSIONS,
+    FUNCTION_CALLS,
+    TRY_PROPAGATION,
+    PREFIX,
+    AS,
+    MULTIPLICATIVE,
+    ADDITIVE,
+    SHIFT,
+    BITWISE_AND,
+    BITWISE_XOR,
+    BITWISE_OR,
+    COMPARISON,
+    LOGICAL_AND,
+    LOGICAL_OR,
+    RANGE,
+    ASSIGN;
+
+    public int value() {
+        return -this.ordinal();
+    }
+}
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriter.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriter.java
index e4b24a8..b584c69 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriter.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriter.java
@@ -294,7 +294,7 @@ public class RustWriter implements AutoCloseable {
     }
 
     private void writeExpressionStatement(RustExpressionStatement expressionStatement) throws IOException {
-        this.writeExpression(expressionStatement.expression());
+        this.writeTopLevelExpression(expressionStatement.expression());
         this.writer.write(";");
     }
 
@@ -307,13 +307,13 @@ public class RustWriter implements AutoCloseable {
 
         this.writeIdentifier(letStatement.variableName());
         this.writer.write(" = ");
-        this.writeExpression(letStatement.value());
+        this.writeTopLevelExpression(letStatement.value());
         this.writer.write(";");
     }
 
     // == Expressions ==
 
-    void writeExpression(RustExpression expression) throws IOException {
+    void writeTopLevelExpression(RustExpression expression) throws IOException {
         switch (expression) {
             case RustArrayRepeatExpression arrayRepeatExpression -> {
                 this.writeArrayRepeatExpression(arrayRepeatExpression);
@@ -369,16 +369,69 @@ public class RustWriter implements AutoCloseable {
         }
     }
 
+    void writeSubExpression(
+        RustExpression expression,
+        RustPrecedence parentPrecedence,
+        boolean matchesAssociativity
+    ) throws IOException {
+        var precedence = precedence(expression);
+        var requiresParens = (
+            parentPrecedence.value() > precedence.value() ||
+                (parentPrecedence.value() == precedence.value() && !matchesAssociativity)
+        );
+
+        if (requiresParens) {
+            this.writer.write("(");
+        }
+
+        this.writeTopLevelExpression(expression);
+
+        if (requiresParens) {
+            this.writer.write(")");
+        }
+    }
+
+    private RustPrecedence precedence(RustExpression expression) {
+        return switch (expression) {
+            case RustArrayRepeatExpression _ -> RustPrecedence.PRIMARY;
+            case RustBinaryExpression binaryExpression -> precedence(binaryExpression.operator());
+            case RustBlockExpression _ -> RustPrecedence.PRIMARY;
+            case RustBoolLiteral _ -> RustPrecedence.PRIMARY;
+            case RustCallExpression _ -> RustPrecedence.FUNCTION_CALLS;
+            case RustFieldExpression _ -> RustPrecedence.FIELD_EXPRESSIONS;
+            case RustIfExpression _ -> RustPrecedence.PRIMARY;
+            case RustIntegerLiteral _ -> RustPrecedence.PRIMARY;
+            case RustPath _ -> RustPrecedence.PATHS;
+            case RustPrefixExpression _ -> RustPrecedence.PREFIX;
+            case RustStructExpression _ -> RustPrecedence.PRIMARY;
+            case RustTryPropagationExpression _ -> RustPrecedence.TRY_PROPAGATION;
+            case RustTupleExpression _ -> RustPrecedence.PRIMARY;
+        };
+    }
+
+    private RustPrecedence precedence(RustBinaryOperator operator) {
+        return switch (operator) {
+            case ADD -> RustPrecedence.ADDITIVE;
+            case EQUAL -> RustPrecedence.COMPARISON;
+            case NOT_EQUAL -> RustPrecedence.COMPARISON;
+        };
+    }
+
     private void writeArrayRepeatExpression(RustArrayRepeatExpression arrayRepeatExpression) throws IOException {
         this.writer.write("[");
-        this.writeExpression(arrayRepeatExpression.repeatOperand());
+        this.writeTopLevelExpression(arrayRepeatExpression.repeatOperand());
         this.writer.write("; ");
-        this.writeExpression(arrayRepeatExpression.lengthOperand());
+        this.writeTopLevelExpression(arrayRepeatExpression.lengthOperand());
         this.writer.write("]");
     }
 
     private void writeBinaryExpression(RustBinaryExpression binaryExpression) throws IOException {
-        this.writeExpression(binaryExpression.left());
+        var associativity = associativity(binaryExpression);
+        this.writeSubExpression(
+            binaryExpression.left(),
+            precedence(binaryExpression),
+            associativity == RustAssociativity.LEFT
+        );
         this.writer.write(" ");
         this.writer.write(switch (binaryExpression.operator()) {
             case ADD -> "+";
@@ -386,7 +439,19 @@ public class RustWriter implements AutoCloseable {
             case NOT_EQUAL -> "!=";
         });
         this.writer.write(" ");
-        this.writeExpression(binaryExpression.right());
+        this.writeSubExpression(
+            binaryExpression.right(),
+            precedence(binaryExpression),
+            associativity == RustAssociativity.RIGHT
+        );
+    }
+
+    private RustAssociativity associativity(RustBinaryExpression binaryExpression) {
+        return switch (binaryExpression.operator()) {
+            case ADD -> RustAssociativity.LEFT;
+            case EQUAL -> RustAssociativity.NONE;
+            case NOT_EQUAL -> RustAssociativity.NONE;
+        };
     }
 
     private void writeBlockExpression(RustBlockExpression blockExpression) throws IOException {
@@ -400,7 +465,7 @@ public class RustWriter implements AutoCloseable {
 
         if (blockExpression.finalOperand().isPresent()) {
             this.writer.newLine();
-            this.writeExpression(blockExpression.finalOperand().get());
+            this.writeTopLevelExpression(blockExpression.finalOperand().get());
         }
 
         this.writer.dedent();
@@ -413,11 +478,11 @@ public class RustWriter implements AutoCloseable {
     }
 
     private void writeCallExpression(RustCallExpression callExpression) throws IOException {
-        this.writeExpression(callExpression.function());
+        this.writeSubExpression(callExpression.function(), precedence(callExpression), true);
         this.writer.write("(");
         writeWithSeparator(
             callExpression.args(),
-            this::writeExpression,
+            this::writeTopLevelExpression,
             () -> {
                 this.writer.write(", ");
             }
@@ -426,14 +491,14 @@ public class RustWriter implements AutoCloseable {
     }
 
     private void writeFieldExpression(RustFieldExpression fieldExpression) throws IOException {
-        this.writeExpression(fieldExpression.containerOperand());
+        this.writeSubExpression(fieldExpression.containerOperand(), precedence(fieldExpression), true);
         this.writer.write(".");
         this.writeIdentifier(fieldExpression.fieldName());
     }
 
     private void writeIfExpression(RustIfExpression ifExpression) throws IOException {
         this.writer.write("if ");
-        this.writeExpression(ifExpression.condition());
+        this.writeTopLevelExpression(ifExpression.condition());
         this.writer.write(" ");
         this.writeBlockExpression(ifExpression.ifTrue());
         this.writer.write(" else ");
@@ -453,7 +518,7 @@ public class RustWriter implements AutoCloseable {
             case BORROW_MUTABLE -> "&mut ";
             case DEREFERENCE -> "*";
         });
-        this.writeExpression(prefixExpression.operand());
+        this.writeSubExpression(prefixExpression.operand(), precedence(prefixExpression), true);
     }
 
     private void writeStructExpression(RustStructExpression structExpression) throws IOException {
@@ -467,7 +532,7 @@ public class RustWriter implements AutoCloseable {
                 fieldValue -> {
                     this.writeIdentifier(fieldValue.name());
                     this.writer.write(": ");
-                    this.writeExpression(fieldValue.value());
+                    this.writeTopLevelExpression(fieldValue.value());
                 },
                 () -> {
                     this.writer.write(", ");
@@ -482,7 +547,11 @@ public class RustWriter implements AutoCloseable {
     private void writeTryPropagationExpression(
         RustTryPropagationExpression tryPropagationExpression
     ) throws IOException {
-        this.writeExpression(tryPropagationExpression.operand());
+        this.writeSubExpression(
+            tryPropagationExpression.operand(),
+            precedence(tryPropagationExpression),
+            true
+        );
         this.writer.write("?");
     }
 
@@ -491,12 +560,12 @@ public class RustWriter implements AutoCloseable {
     ) throws IOException {
         this.writer.write("(");
         if (tupleExpression.elements().size() == 1) {
-            this.writeExpression(tupleExpression.elements().getFirst());
+            this.writeTopLevelExpression(tupleExpression.elements().getFirst());
             this.writer.write(",");
         } else if (tupleExpression.elements().size() > 1) {
             writeWithSeparator(
                 tupleExpression.elements(),
-                this::writeExpression,
+                this::writeTopLevelExpression,
                 () -> {
                     this.writer.write(", ");
                 }