From cf52b6d9b8e009a1fa8112b01a5158ee0f82cf0e Mon Sep 17 00:00:00 2001
From: Madhanika Jayasimha <madhanika.jayasimha@intel.com>
Date: Tue, 18 Jul 2023 15:00:47 -0700
Subject: [PATCH] Added changes based on feedback

---
 .../signal_width_mismatch_exception.dart      |  2 +-
 lib/src/modules/conditional.dart              | 52 ++++++++++++++-----
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/lib/src/exceptions/logic/signal_width_mismatch_exception.dart b/lib/src/exceptions/logic/signal_width_mismatch_exception.dart
index 21eda4ab4..895cb6510 100644
--- a/lib/src/exceptions/logic/signal_width_mismatch_exception.dart
+++ b/lib/src/exceptions/logic/signal_width_mismatch_exception.dart
@@ -24,6 +24,6 @@ class SignalWidthMismatchException extends RohdException {
       dynamic val, int expectedWidth, int actualWidth,
       {String additionalMessage = ''})
       : super('Value $val has the wrong width.'
-            ' Excepted $expectedWidth but found $actualWidth.'
+            ' Expected $expectedWidth but found $actualWidth.'
             ' $additionalMessage');
 }
diff --git a/lib/src/modules/conditional.dart b/lib/src/modules/conditional.dart
index 6749868b6..d24268ca4 100644
--- a/lib/src/modules/conditional.dart
+++ b/lib/src/modules/conditional.dart
@@ -889,38 +889,66 @@ enum ConditionalType {
   priority
 }
 
-/// Shorthand for a [Conditional] inside a [Case] block.
-/// Returns a [Logic] signal based on the expression and conditions.
-/// The width of the result is determined by the width of the value of the
-/// conditions.
-Logic cases(Logic expression, Map<Logic, dynamic> conditions, {int? width}) {
+/// Shorthand for a [Case] inside a [Conditional] block.
+///
+/// It is used to assign a signal based on a condition with multiple cases to
+/// consider. For e.g., this can be used instead of a nested [mux].
+///
+/// The result is of type [Logic] and it is determined by conditionaly matching
+/// the expression with the values of each item in conditions. If width of the
+/// input is not provided, then the width of  the result is inferred from the
+/// width of the entries.
+Logic cases(Logic expression, Map<dynamic, dynamic> conditions,
+    {int? width, ConditionalType conditionalType = ConditionalType.none}) {
   for (final condition in conditions.entries) {
     int? inferredWidth;
+
     if (condition.value is Logic) {
       inferredWidth = (condition.value as Logic).width;
     } else if (condition.value is LogicValue) {
       inferredWidth = (condition.value as LogicValue).width;
     }
 
+    if (width != inferredWidth && width != null && inferredWidth != null) {
+      throw SignalWidthMismatchException.forDynamic(
+          condition.value, width, inferredWidth);
+    }
+
     width ??= inferredWidth;
 
     if (width != inferredWidth && inferredWidth != null) {
       throw SignalWidthMismatchException.forDynamic(
           condition.value, width!, inferredWidth);
     }
-    if (expression.width != condition.key.width) {
-      throw SignalWidthMismatchException.forDynamic(
-          condition.key, expression.width, condition.key.width);
+    if (condition.key is Logic) {
+      if (expression.width != (condition.key as Logic).width) {
+        throw SignalWidthMismatchException.forDynamic(
+            condition.key, expression.width, (condition.key as Logic).width);
+      }
+    }
+
+    if (condition.key is LogicValue) {
+      if (expression.width != (condition.key as LogicValue).width) {
+        throw SignalWidthMismatchException.forDynamic(condition.key,
+            expression.width, (condition.key as LogicValue).width);
+      }
     }
   }
 
   final result = Logic(name: 'result', width: width!);
 
   Combinational([
-    Case(expression, [
-      for (final condition in conditions.entries)
-        CaseItem(condition.key, [result < condition.value])
-    ])
+    Case(
+        expression,
+        [
+          for (final condition in conditions.entries)
+            CaseItem(
+                condition.key is Logic
+                    ? condition.key as Logic
+                    : Const(condition.key, width: expression.width),
+                [result < condition.value])
+        ],
+        conditionalType: conditionalType)
   ]);
 
   return result;