OneHotEnum: a ChiselEnum with one-hot encoding

core/src/main/scala-2/chisel3/OneHotEnumMacros.scala

@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package chisel3
+
+import scala.reflect.macros.blackbox.Context
+import scala.language.experimental.macros
+
+private[chisel3] trait OneHotEnumIntf extends ChiselEnumIntf { self: OneHotEnum =>
+  override def Value:           Type = macro OneHotEnumMacros.ValImpl
+  override def Value(id: UInt): Type = macro OneHotEnumMacros.ValCustomImpl
+}
+
+private[chisel3] object OneHotEnumMacros {
+  def ValImpl(c: Context): c.Tree = {
+    import c.universe._
+
+    val term = c.internal.enclosingOwner
+    val name = term.name.decodedName.toString.trim
+
+    if (name.contains(" ")) {
+      c.abort(c.enclosingPosition, "Value cannot be called without assigning to an enum")
+    }
+
+    q"""this.do_OHValue($name)"""
+  }
+
+  def ValCustomImpl(c: Context)(id: c.Expr[UInt]): c.universe.Tree = {
+    c.abort(c.enclosingPosition, "OneHotEnum does not support custom values")
+  }
+}

core/src/main/scala-3/chisel3/ChiselEnumIntf.scala

@@ -17,3 +17,7 @@ private[chisel3] trait EnumTypeIntf { self: EnumType =>
 private[chisel3] trait ChiselEnumIntf { self: ChiselEnum =>
   // TODO macros
 }
+
+private[chisel3] trait OneHotEnumIntf extends ChiselEnumIntf { self: OneHotEnum =>
+  // TODO macros
+}

core/src/main/scala/chisel3/ChiselEnum.scala

@@ -92,47 +92,11 @@ abstract class EnumType(private[chisel3] val factory: ChiselEnum) extends Elemen
     if (litOption.isDefined) {
       true.B
     } else {
-      if (factory.isTotal) true.B else factory.all.map(this === _).reduce(_ || _)
+      if (factory.isTotal) true.B else factory._isValid(this)
     }
   }
 
-  /** Test if this enumeration is equal to any of the values in a given sequence
-    *
-    * @param s a [[scala.collection.Seq$ Seq]] of enumeration values to look for
-    * @return a hardware [[Bool]] that indicates if this value matches any of the given values
-    */
-  final def isOneOf(s: Seq[EnumType])(implicit sourceInfo: SourceInfo): Bool = {
-    VecInit(s.map(this === _)).asUInt.orR
-  }
-
-  /** Test if this enumeration is equal to any of the values given as arguments
-    *
-    * @param u1 the first value to look for
-    * @param u2 zero or more additional values to look for
-    * @return a hardware [[Bool]] that indicates if this value matches any of the given values
-    */
-  final def isOneOf(
-    u1: EnumType,
-    u2: EnumType*
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Bool = isOneOf(u1 +: u2.toSeq)
-
-  def next(implicit sourceInfo: SourceInfo): this.type = {
-    if (litOption.isDefined) {
-      val index = factory.all.indexOf(this)
-
-      if (index < factory.all.length - 1) {
-        factory.all(index + 1).asInstanceOf[this.type]
-      } else {
-        factory.all.head.asInstanceOf[this.type]
-      }
-    } else {
-      val enums_with_nexts = factory.all.zip(factory.all.tail :+ factory.all.head)
-      val next_enum = SeqUtils.priorityMux(enums_with_nexts.map { case (e, n) => (this === e, n) })
-      next_enum.asInstanceOf[this.type]
-    }
-  }
+  def next(implicit sourceInfo: SourceInfo): this.type = factory._next(this)
 
   private[chisel3] def bindToLiteral(num: BigInt, w: Width): Unit = {
     val lit = ULit(num, w)
@@ -203,7 +167,59 @@ abstract class EnumType(private[chisel3] val factory: ChiselEnum) extends Elemen
 }
 
 abstract class ChiselEnum extends ChiselEnumIntf {
-  class Type extends EnumType(this)
+  private[chisel3] protected def _valIs(v: Type, lit: Type)(implicit sourceInfo: SourceInfo): Bool = {
+    v === lit
+  }
+
+  private[chisel3] protected def _valIsOneOf(v: Type, s: Seq[Type])(implicit sourceInfo: SourceInfo): Bool =
+    VecInit(s.map(_valIs(v, _))).asUInt.orR
+
+  private[chisel3] protected def _isValid(v: EnumType)(implicit sourceInfo: SourceInfo): Bool = {
+    assert(v.isInstanceOf[Type])
+    all.map(v === _).reduce(_ || _)
+  }
+
+  private[chisel3] protected def _next(v: EnumType)(implicit sourceInfo: SourceInfo): v.type = {
+    assert(v.isInstanceOf[Type])
+
+    if (v.litOption.isDefined) {
+      val index = v.factory.all.indexOf(v)
+
+      if (index < v.factory.all.length - 1) {
+        v.factory.all(index + 1).asInstanceOf[v.type]
+      } else {
+        v.factory.all.head.asInstanceOf[v.type]
+      }
+    } else {
+      val enums_with_nexts = v.factory.all.zip(v.factory.all.tail :+ v.factory.all.head)
+      val next_enum = SeqUtils.priorityMux(enums_with_nexts.map { case (e, n) => (v === e, n) })
+      next_enum.asInstanceOf[v.type]
+    }
+  }
+
+  class Type extends EnumType(this) {
+
+    /** Test if this enumeration is equal to any of the values in a given sequence
+      *
+      * @param s a [[scala.collection.Seq$ Seq]] of enumeration values to look for
+      * @return a hardware [[Bool]] that indicates if this value matches any of the given values
+      */
+    final def isOneOf(s: Seq[Type])(implicit sourceInfo: SourceInfo): Bool = _valIsOneOf(this, s)
+
+    /** Test if this enumeration is equal to any of the values given as arguments
+      *
+      * @param u1 the first value to look for
+      * @param u2 zero or more additional values to look for
+      * @return a hardware [[Bool]] that indicates if this value matches any of the given values
+      */
+    final def isOneOf(
+      u1: Type,
+      u2: Type*
+    )(
+      implicit sourceInfo: SourceInfo
+    ): Bool = isOneOf(u1 +: u2.toSeq)
+  }
+
   object Type {
     def apply(): Type = ChiselEnum.this.apply()
   }

core/src/main/scala/chisel3/OneHotEnum.scala

@@ -0,0 +1,89 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package chisel3
+
+import chisel3.experimental.SourceInfo
+
+abstract class OneHotEnum extends ChiselEnum with OneHotEnumIntf {
+  private var next1Pos = 0
+
+  // copied from chisel3.util
+  private def isPow2(in:   BigInt): Boolean = in > 0 && ((in & (in - 1)) == 0)
+  private def log2Ceil(in: BigInt): Int = (in - 1).bitLength
+
+  override def _valIs(v: Type, lit: Type)(implicit sourceInfo: SourceInfo): Bool = {
+    require(lit.isLit, "Can only compare against literal values")
+    val yy = lit.litValue
+    require(isPow2(yy), s"Can only compare against one-hot values, got $yy (0b${yy.toString(2)})")
+    v.asUInt.apply(log2Ceil(yy))
+  }
+
+  override def _isValid(v: EnumType)(implicit sourceInfo: SourceInfo): Bool = {
+    assert(v.isInstanceOf[Type])
+    assert(v.getWidth == all.length, s"OneHotEnum ${this} has ${all.length} values, but ${v} has width ${v.getWidth}")
+    val x = v.asUInt
+    x.orR && ((x & (x - 1.U)) === 0.U)
+  }
+
+  override def _next(v: EnumType)(implicit sourceInfo: SourceInfo): v.type = {
+    assert(v.isInstanceOf[Type])
+
+    if (v.litOption.isDefined) {
+      val index = v.factory.all.indexOf(v)
+
+      if (index < v.factory.all.length - 1) {
+        v.factory.all(index + 1).asInstanceOf[v.type]
+      } else {
+        v.factory.all.head.asInstanceOf[v.type]
+      }
+    } else {
+      safe(v.asUInt.rotateLeft(1))._1.asInstanceOf[v.type]
+    }
+  }
+
+  override def isTotal: Boolean = false
+
+  // TODO: Is there a cleaner way?
+  final implicit class OneHotType(value: Type) extends Type {
+    override def isLit: Boolean = value.isLit
+
+    override def litValue: BigInt = value.litValue
+
+    final def is(other: Type)(implicit sourceInfo: SourceInfo): Bool = _valIs(value, other)
+
+    /**
+      * Multiplexer that selects between multiple values based on this one-hot enum.
+      *
+      * @param choices a sequence of tuples of (enum value, output when matched)
+      * @return the output corresponding to the matched enum value
+      * @note the output is undefined if none of the values match 
+      */
+    final def select[T <: Data](choices: Seq[(Type, T)])(implicit sourceInfo: SourceInfo): T = {
+      require(choices.nonEmpty, "select must be passed a non-empty list of choices")
+      // FIXME: this is a workaround to suppress a superfluous cast warning emitted by [[SeqUtils.oneHotMux]] when T is of the same EnumType. Unfortunately, it also hides potential cast warnings from the inner expressions.
+      suppressEnumCastWarning {
+        SeqUtils.oneHotMux(choices.map { case (oh, t) => is(oh) -> t })
+      }
+    }
+
+    /**
+      * Multiplexer that selects between multiple values based on this one-hot enum.
+      *
+      * @param firstChoice a tuple of (enum value, output when matched)
+      * @param otherChoices a varargs list of tuples of (enum value, output when matched)
+      * @return the output corresponding to the matched enum value
+      * @note if none of the enum values match, the output is undefined
+      */
+    final def select[T <: Data](
+      firstChoice:  (Type, T),
+      otherChoices: (Type, T)*
+    )(implicit sourceInfo: SourceInfo): T = select(firstChoice +: otherChoices)
+
+  }
+
+  def do_OHValue(name: String): Type = {
+    val value = super.do_Value(name, BigInt(2).pow(next1Pos).U)
+    next1Pos += 1
+    value
+  }
+}