Extract generateEncodeInt() and generateDecodeInt()

1 files changed, 30 insertions, 5 deletions

diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatransient0/JavaTransient0Generator.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatransient0/JavaTransient0Generator.java
index ee35ebc..02b2a3f 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatransient0/JavaTransient0Generator.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatransient0/JavaTransient0Generator.java
@@ -18,6 +18,7 @@ import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
 public class JavaTransient0Generator implements Generator {
@@ -250,8 +251,21 @@ public class JavaTransient0Generator implements Generator {
     }
 
     private static Stream<JavaBlockStatement> generateEncodeInt32(JavaExpression value) {
-        return Stream.of(0, 8, 16, 24)
-            .map(shiftBy -> new JavaExpressionStatement(
+        return generateEncodeInt(value, 32);
+    }
+
+    private static Stream<JavaBlockStatement> generateDecodeInt32(
+        JavaIdentifier target
+    ) {
+        return generateDecodeInt(target, 32);
+    }
+
+    private static Stream<JavaBlockStatement> generateEncodeInt(
+        JavaExpression value,
+        int bits
+    ) {
+        return byteIndicesInBits(bits)
+            .mapToObj(shiftBy -> new JavaExpressionStatement(
                 new JavaMethodCall(
                     new JavaRef(OUTPUT_STREAM_NAME),
                     JavaIdentifier.of("write"),
@@ -268,14 +282,18 @@ public class JavaTransient0Generator implements Generator {
             ));
     }
 
-    private static Stream<JavaBlockStatement> generateDecodeInt32(
-        JavaIdentifier target
+    private static Stream<JavaBlockStatement> generateDecodeInt(
+        JavaIdentifier target,
+        int bits
     ) {
         var statements = new ArrayList<JavaBlockStatement>();
 
         statements.add(new JavaLocalVariableDeclaration(target, new JavaIntegerLiteral(0)));
 
-        for (var shiftBy : List.of(0, 8, 16, 24)) {
+        Iterable<Integer> shiftBys = () -> byteIndicesInBits(bits)
+            .iterator();
+
+        for (var shiftBy : shiftBys) {
             statements.add(new JavaExpressionStatement(new JavaBinaryOperation(
                 JavaBinaryOperator.ASSIGN,
                 new JavaRef(target),
@@ -298,6 +316,13 @@ public class JavaTransient0Generator implements Generator {
         return statements.stream();
     }
 
+    /// For an integer of the given number of bits, find the indices of each
+    /// byte in bits.
+    private static IntStream byteIndicesInBits(int bits) {
+        return IntStream.range(0, bits / 8)
+            .map(x -> x * 8);
+    }
+
     private JavaIdentifier encodeMethodName(StructType type) {
         return this.javaGenerator.generateFieldName("encode" + type.name());
     }