Implement more precise calculation of size for switch-on types

kaitai-io · Sep 12, 2023 · 527ba04 · 527ba04
1 parent 67ad00a
commit 527ba04
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 3 deletions.
diff --git a/jvm/src/test/scala/io/kaitai/struct/CalculateSeqSizes$Test.scala b/jvm/src/test/scala/io/kaitai/struct/CalculateSeqSizes$Test.scala
@@ -3,7 +3,7 @@ package io.kaitai.struct
 import io.kaitai.struct.datatype.{BigEndian, BigBitEndian}
 import io.kaitai.struct.datatype.DataType
 import io.kaitai.struct.datatype.DataType._
-import io.kaitai.struct.exprlang.Ast
+import io.kaitai.struct.exprlang.{Ast, Expressions}
 import io.kaitai.struct.format.{DynamicSized, FixedSized, MetaSpec, Sized, YamlAttrArgs}
 import io.kaitai.struct.precompile.CalculateSeqSizes
 
@@ -72,6 +72,16 @@ class CalculateSeqSizes$Test extends FunSpec {
     sizeof(None, None, None, Some(contents))
   }
 
+  /** Helper for testing the `switch-on` types. */
+  private def switchOn(cases: Map[String, String]): Sized = {
+    CalculateSeqSizes.dataTypeBitsSize(SwitchType(
+      Ast.expr.IntNum(0),
+      cases.map { case (condition, typeName) =>
+        Expressions.parse(condition) -> parse(Some(typeName), None, None, None)
+      }
+    ))
+  }
+
   describe("CalculateSeqSizes") {
     it("built-in types has correct size") {
       sizeof("s1") should be (FixedSized( 8))
@@ -155,5 +165,49 @@ class CalculateSeqSizes$Test extends FunSpec {
 
       sizeof("abcdef".getBytes()) should be (FixedSized(6*8))// content: 'abcdef'
     }
+
+    describe("switch-on") {
+      it("has a zero size if no cases present") {
+        switchOn(Map()) should be (FixedSized(0))
+      }
+
+      it("has a fixed size when all cases have the same fixed size") {
+        switchOn(Map(
+          "0" -> "f4be",
+          "1" -> "u4be",
+          "_" -> "s4be",
+        )) should be (FixedSized(4*8))
+      }
+
+      it("has a dynamic size when not all cases have the same size") {
+        switchOn(Map(
+          "0" -> "f4be",
+          "1" -> "u4be",
+          "_" -> "u1",
+        )) should be (DynamicSized)
+      }
+
+      it("has a dynamic size when contains a case with a dynamic size") {
+        // Fixed + Dynamic
+        switchOn(Map(
+          "0" -> "f4be",
+          "1" -> "u4be",
+          "_" -> "strz",
+        )) should be (DynamicSized)
+
+        // Dynamic + Fixed
+        switchOn(Map(
+          "0" -> "strz",
+          "1" -> "u4be",
+          "_" -> "f4be",
+        )) should be (DynamicSized)
+
+        // Dynamic + Dynamic
+        switchOn(Map(
+          "1" -> "strz",
+          "_" -> "strz",
+        )) should be (DynamicSized)
+      }
+    }
   }
 }
diff --git a/shared/src/main/scala/io/kaitai/struct/format/ClassSpec.scala b/shared/src/main/scala/io/kaitai/struct/format/ClassSpec.scala
@@ -25,7 +25,26 @@ case object GenericStructClassSpec extends ClassSpecLike {
   * Type that represent result of the `_sizeof` special property and `sizeof<>`
   * meta-function.
   */
-sealed trait Sized
+sealed trait Sized {
+  /**
+    * Combines two sizes into one object which covers both size requirements
+    *
+    * @param other Size to combine with
+    */
+  final def or(other: Sized): Sized = {
+    (this, other) match {
+      case (FixedSized(l), FixedSized(r)) if l == r => FixedSized(l)
+      case (FixedSized(_), FixedSized(_)) => DynamicSized
+
+      case (FixedSized(_), DynamicSized) => DynamicSized
+      case (DynamicSized, FixedSized(_)) => DynamicSized
+
+      case (DynamicSized, DynamicSized) => DynamicSized
+
+      // other combinations should produce an error
+    }
+  }
+}
 /**
   * The size of type have no constant value. The examples is built-in unsized
   * types: `str`, `strz`, and `bytes`. Those types has no natural size in contrary

diff --git a/shared/src/main/scala/io/kaitai/struct/precompile/CalculateSeqSizes.scala b/shared/src/main/scala/io/kaitai/struct/precompile/CalculateSeqSizes.scala
@@ -119,7 +119,12 @@ object CalculateSeqSizes {
       case StrFromBytesType(basedOn, _) => dataTypeBitsSize(basedOn)
       case CalcStrType => DynamicSized
 
-      case st: SwitchType => DynamicSized // FIXME: it's really possible get size if st.hasSize
+      case SwitchType(_, cases, _) => cases.values.foldLeft[Option[Sized]](None) {
+        case (acc, dataType) => Some(acc match {
+          case None    => dataTypeBitsSize(dataType)
+          case Some(s) => s.or(dataTypeBitsSize(dataType))
+        })
+      }.getOrElse(FixedSized(0))
 
       case OwnedKaitaiStreamType | KaitaiStreamType => DynamicSized