@@ -20,8 +20,15 @@ object DesugarEnums {
2020 val Simple, Object, Class : Value = Value
2121 }
2222
23- /** Attachment containing the number of enum cases and the smallest kind that was seen so far. */
24- val EnumCaseCount : Property .Key [(Int , DesugarEnums .CaseKind .Value )] = Property .Key ()
23+ /** Attachment containing the number of enum cases, the smallest kind that was seen so far,
24+ * and a list of all the value cases with their ordinals.
25+ */
26+ val EnumCaseCount : Property .Key [(Int , CaseKind .Value , List [(Int , TermName )])] = Property .Key ()
27+
28+ /** Attachment signalling that when this definition is desugared, it should add any additional
29+ * lookup methods for enums.
30+ */
31+ val DefinesEnumLookupMethods : Property .Key [Unit ] = Property .Key ()
2532
2633 /** The enumeration class that belongs to an enum case. This works no matter
2734 * whether the case is still in the enum class or it has been transferred to the
@@ -122,6 +129,21 @@ object DesugarEnums {
122129 valueOfDef :: Nil
123130 }
124131
132+ private def enumLookupMethods (cases : List [(Int , TermName )])(using Context ): List [Tree ] =
133+ if isJavaEnum || cases.isEmpty then Nil
134+ else
135+ val defaultCase =
136+ val ord = Ident (nme.ordinal)
137+ val err = Throw (New (TypeTree (defn.IndexOutOfBoundsException .typeRef), List (Select (ord, nme.toString_) :: Nil )))
138+ CaseDef (ord, EmptyTree , err)
139+ val valueCases = cases.map((i, name) =>
140+ CaseDef (Literal (Constant (i)), EmptyTree , Ident (name))
141+ ) ::: defaultCase :: Nil
142+ val fromOrdinalDef = DefDef (nme.fromOrdinalDollar, Nil , List (param(nme.ordinalDollar_, defn.IntType ) :: Nil ),
143+ rawRef(enumClass.typeRef), Match (Ident (nme.ordinalDollar_), valueCases))
144+ .withFlags(Synthetic | Private )
145+ fromOrdinalDef :: Nil
146+
125147 /** A creation method for a value of enum type `E`, which is defined as follows:
126148 *
127149 * private def $new(_$ordinal: Int, $name: String) = new E with scala.runtime.EnumValue {
@@ -256,16 +278,22 @@ object DesugarEnums {
256278 * - scaffolding containing the necessary definitions for singleton enum cases
257279 * unless that scaffolding was already generated by a previous call to `nextEnumKind`.
258280 */
259- def nextOrdinal (kind : CaseKind .Value )(using Context ): (Int , List [Tree ]) = {
260- val (count, seenKind) = ctx.tree.removeAttachment(EnumCaseCount ).getOrElse((0 , CaseKind .Class ))
261- val minKind = if (kind < seenKind) kind else seenKind
262- ctx.tree.pushAttachment(EnumCaseCount , (count + 1 , minKind))
263- val scaffolding =
281+ def nextOrdinal (name : Name , kind : CaseKind .Value , definesLookups : Boolean )(using Context ): (Int , List [Tree ]) = {
282+ val (ordinal, seenKind, seenCases) = ctx.tree.removeAttachment(EnumCaseCount ).getOrElse((0 , CaseKind .Class , Nil ))
283+ val minKind = if kind < seenKind then kind else seenKind
284+ val cases = name match
285+ case name : TermName => (ordinal, name) :: seenCases
286+ case _ => seenCases
287+ ctx.tree.pushAttachment(EnumCaseCount , (ordinal + 1 , minKind, cases))
288+ val scaffolding0 =
264289 if (kind >= seenKind) Nil
265290 else if (kind == CaseKind .Object ) enumScaffolding
266291 else if (seenKind == CaseKind .Object ) enumValueCreator :: Nil
267292 else enumScaffolding :+ enumValueCreator
268- (count, scaffolding)
293+ val scaffolding =
294+ if definesLookups then scaffolding0 ::: enumLookupMethods(cases.reverse)
295+ else scaffolding0
296+ (ordinal, scaffolding)
269297 }
270298
271299 def param (name : TermName , typ : Type )(using Context ) =
@@ -286,13 +314,13 @@ object DesugarEnums {
286314 enumLabelMeth(Literal (Constant (name)))
287315
288316 /** Expand a module definition representing a parameterless enum case */
289- def expandEnumModule (name : TermName , impl : Template , mods : Modifiers , span : Span )(using Context ): Tree = {
317+ def expandEnumModule (name : TermName , impl : Template , mods : Modifiers , definesLookups : Boolean , span : Span )(using Context ): Tree = {
290318 assert(impl.body.isEmpty)
291319 if (! enumClass.exists) EmptyTree
292320 else if (impl.parents.isEmpty)
293- expandSimpleEnumCase(name, mods, span)
321+ expandSimpleEnumCase(name, mods, definesLookups, span)
294322 else {
295- val (tag, scaffolding) = nextOrdinal(CaseKind .Object )
323+ val (tag, scaffolding) = nextOrdinal(name, CaseKind .Object , definesLookups )
296324 val ordinalDef = if isJavaEnum then Nil else ordinalMethLit(tag) :: Nil
297325 val enumLabelDef = enumLabelLit(name.toString)
298326 val impl1 = cpy.Template (impl)(
@@ -305,15 +333,15 @@ object DesugarEnums {
305333 }
306334
307335 /** Expand a simple enum case */
308- def expandSimpleEnumCase (name : TermName , mods : Modifiers , span : Span )(using Context ): Tree =
336+ def expandSimpleEnumCase (name : TermName , mods : Modifiers , definesLookups : Boolean , span : Span )(using Context ): Tree =
309337 if (! enumClass.exists) EmptyTree
310338 else if (enumClass.typeParams.nonEmpty) {
311339 val parent = interpolatedEnumParent(span)
312340 val impl = Template (emptyConstructor, parent :: Nil , Nil , EmptyValDef , Nil )
313- expandEnumModule(name, impl, mods, span)
341+ expandEnumModule(name, impl, mods, definesLookups, span)
314342 }
315343 else {
316- val (tag, scaffolding) = nextOrdinal(CaseKind .Simple )
344+ val (tag, scaffolding) = nextOrdinal(name, CaseKind .Simple , definesLookups )
317345 val creator = Apply (Ident (nme.DOLLAR_NEW ), List (Literal (Constant (tag)), Literal (Constant (name.toString))))
318346 val vdef = ValDef (name, enumClassRef, creator).withMods(mods.withAddedFlags(EnumValue , span))
319347 flatTree(scaffolding ::: vdef :: Nil ).withSpan(span)
0 commit comments