jsonutils: add customization for toJson via `ToJsonOptions`; generalize symbolName; add symbolRank (#18029)

* jsonutils: add customization for toJson via `ToJsonOptions`

* add enumutils.symbolRank

* lookup table implementation for HoleyEnum

* cleanup

* changelog

* fixup

* Update lib/std/jsonutils.nim

Co-authored-by: Andreas Rumpf <rumpf_a@web.de>

4 files changed, 112 insertions, 14 deletions

diff --git a/changelog.md b/changelog.md
index 369b11917..97e749337 100644
--- a/changelog.md
+++ b/changelog.md
@@ -118,9 +118,11 @@
 - `json.%`,`json.to`, `jsonutils.formJson`,`jsonutils.toJson` now work with `uint|uint64`
   instead of raising (as in 1.4) or giving wrong results (as in 1.2).
 
+- `jsonutils` now handles `cstring` (including as Table key), and `set`.
+
 - added `jsonutils.jsonTo` overload with `opt = Joptions()` param.
 
-- `jsonutils` now handles `cstring` (including as Table key), and `set`.
+- `jsonutils.toJson` now supports customization via `ToJsonOptions`.
 
 - Added an overload for the `collect` macro that inferes the container type based
   on the syntax of the last expression. Works with std seqs, tables and sets.
@@ -138,6 +140,7 @@
 - Added `std/enumutils` module. Added `genEnumCaseStmt` macro that generates case statement to parse string to enum.
   Added `items` for enums with holes.
   Added `symbolName` to return the enum symbol name ignoring the human readable name.
+  Added `symbolRank` to return the index in which an enum member is listed in an enum.
 
 - Added `typetraits.HoleyEnum` for enums with holes, `OrdinalEnum` for enums without holes.
 
diff --git a/lib/std/enumutils.nim b/lib/std/enumutils.nim
index 6195ae07d..09cf24f51 100644
--- a/lib/std/enumutils.nim
+++ b/lib/std/enumutils.nim
@@ -86,8 +86,67 @@ iterator items*[T: HoleyEnum](E: typedesc[T]): T =
     assert B[float].toSeq == [B[float].b0, B[float].b1]
   for a in enumFullRange(E): yield a
 
-func symbolName*[T: OrdinalEnum](a: T): string =
+func span(T: typedesc[HoleyEnum]): int =
+  (T.high.ord - T.low.ord) + 1
+
+const invalidSlot = uint8.high
+
+proc genLookup[T: typedesc[HoleyEnum]](_: T): auto =
+  const n = span(T)
+  var ret: array[n, uint8]
+  var i = 0
+  assert n <= invalidSlot.int
+  for ai in mitems(ret): ai = invalidSlot
+  for ai in items(T):
+    ret[ai.ord - T.low.ord] = uint8(i)
+    inc(i)
+  return ret
+
+func symbolRankImpl[T](a: T): int {.inline.} =
+  const n = T.span
+  const thres = 255 # must be <= `invalidSlot`, but this should be tuned.
+  when n <= thres:
+    const lookup = genLookup(T)
+    let lookup2 {.global.} = lookup # xxx improve pending https://github.com/timotheecour/Nim/issues/553
+    #[
+    This could be optimized using a hash adapted to `T` (possible since it's known at CT)
+    to get better key distribution before indexing into the lookup table table.
+    ]#
+    {.noSideEffect.}: # because it's immutable
+      let ret = lookup2[ord(a) - T.low.ord]
+    if ret != invalidSlot: return ret.int
+  else:
+    var i = 0
+    # we could also generate a case statement as optimization
+    for ai in items(T):
+      if ai == a: return i
+      inc(i)
+  raise newException(IndexDefect, $ord(a) & " invalid for " & $T)
+
+template symbolRank*[T: enum](a: T): int =
+  ## Returns the index in which `a` is listed in `T`.
+  ##
+  ## The cost for a `HoleyEnum` is implementation defined, currently optimized
+  ## for small enums, otherwise is `O(T.enumLen)`.
+  runnableExamples:
+    type
+      A = enum a0 = -3, a1 = 10, a2, a3 = (20, "f3Alt") # HoleyEnum
+      B = enum b0, b1, b2 # OrdinalEnum
+      C = enum c0 = 10, c1, c2 # OrdinalEnum
+    assert a2.symbolRank == 2
+    assert b2.symbolRank == 2
+    assert c2.symbolRank == 2
+    assert c2.ord == 12
+    assert a2.ord == 11
+    var invalid = 7.A
+    doAssertRaises(IndexDefect): discard invalid.symbolRank
+  when T is Ordinal: ord(a) - T.low.ord.static
+  else: symbolRankImpl(a)
+
+func symbolName*[T: enum](a: T): string =
   ## Returns the symbol name of an enum.
+  ##
+  ## This uses `symbolRank`.
   runnableExamples:
     type B = enum
       b0 = (10, "kb0")
@@ -97,5 +156,7 @@ func symbolName*[T: OrdinalEnum](a: T): string =
     assert b.symbolName == "b0"
     assert $b == "kb0"
     static: assert B.high.symbolName == "b2"
+    type C = enum c0 = -3, c1 = 4, c2 = 20 # HoleyEnum
+    assert c1.symbolName == "c1"
   const names = enumNames(T)
-  names[a.ord - T.low.ord]
+  names[a.symbolRank]
diff --git a/lib/std/jsonutils.nim b/lib/std/jsonutils.nim
index 1f49f60ed..1e222e3a2 100644
--- a/lib/std/jsonutils.nim
+++ b/lib/std/jsonutils.nim
@@ -31,9 +31,11 @@ add a way to customize serialization, for e.g.:
 ]#
 
 import macros
+from enumutils import symbolName
+from typetraits import OrdinalEnum
 
 type
-  Joptions* = object
+  Joptions* = object # xxx rename FromJsonOptions
     ## Options controlling the behavior of `fromJson`.
     allowExtraKeys*: bool
       ## If `true` Nim's object to which the JSON is parsed is not required to
@@ -42,6 +44,17 @@ type
       ## If `true` Nim's object to which JSON is parsed is allowed to have
       ## fields without corresponding JSON keys.
     # in future work: a key rename could be added
+  EnumMode* = enum
+    joptEnumOrd
+    joptEnumSymbol
+    joptEnumString
+  ToJsonOptions* = object
+    enumMode*: EnumMode
+    # xxx charMode
+
+proc initToJsonOptions*(): ToJsonOptions =
+  ## initializes `ToJsonOptions` with sane options.
+  ToJsonOptions(enumMode: joptEnumOrd)
 
 proc isNamedTuple(T: typedesc): bool {.magic: "TypeTrait".}
 proc distinctBase(T: typedesc): typedesc {.magic: "TypeTrait".}
@@ -261,33 +274,41 @@ proc jsonTo*(b: JsonNode, T: typedesc, opt = Joptions()): T =
   ## reverse of `toJson`
   fromJson(result, b, opt)
 
-proc toJson*[T](a: T): JsonNode =
+proc toJson*[T](a: T, opt = initToJsonOptions()): JsonNode =
   ## serializes `a` to json; uses `toJsonHook(a: T)` if it's in scope to
   ## customize serialization, see strtabs.toJsonHook for an example.
   when compiles(toJsonHook(a)): result = toJsonHook(a)
   elif T is object | tuple:
     when T is object or isNamedTuple(T):
       result = newJObject()
-      for k, v in a.fieldPairs: result[k] = toJson(v)
+      for k, v in a.fieldPairs: result[k] = toJson(v, opt)
     else:
       result = newJArray()
-      for v in a.fields: result.add toJson(v)
+      for v in a.fields: result.add toJson(v, opt)
   elif T is ref | ptr:
     if system.`==`(a, nil): result = newJNull()
-    else: result = toJson(a[])
+    else: result = toJson(a[], opt)
   elif T is array | seq | set:
     result = newJArray()
-    for ai in a: result.add toJson(ai)
-  elif T is pointer: result = toJson(cast[int](a))
+    for ai in a: result.add toJson(ai, opt)
+  elif T is pointer: result = toJson(cast[int](a), opt)
     # edge case: `a == nil` could've also led to `newJNull()`, but this results
     # in simpler code for `toJson` and `fromJson`.
-  elif T is distinct: result = toJson(a.distinctBase)
+  elif T is distinct: result = toJson(a.distinctBase, opt)
   elif T is bool: result = %(a)
   elif T is SomeInteger: result = %a
-  elif T is Ordinal: result = %(a.ord)
   elif T is enum:
-    when defined(nimLegacyJsonutilsHoleyEnum): result = %a
-    else: result = %(a.ord)
+    case opt.enumMode
+    of joptEnumOrd:
+      when T is Ordinal or not defined(nimLegacyJsonutilsHoleyEnum): %(a.ord)
+      else: toJson($a, opt)
+    of joptEnumSymbol:
+      when T is OrdinalEnum:
+        toJson(symbolName(a), opt)
+      else:
+        toJson($a, opt)
+    of joptEnumString: toJson($a, opt)
+  elif T is Ordinal: result = %(a.ord)
   elif T is cstring: (if a == nil: result = newJNull() else: result = % $a)
   else: result = %a
 
diff --git a/tests/stdlib/tjsonutils.nim b/tests/stdlib/tjsonutils.nim
index c826a79b0..a55b0ca1d 100644
--- a/tests/stdlib/tjsonutils.nim
+++ b/tests/stdlib/tjsonutils.nim
@@ -35,6 +35,13 @@ type Foo = ref object
 proc `==`(a, b: Foo): bool =
   a.id == b.id
 
+type MyEnum = enum me0, me1 = "me1Alt", me2, me3, me4
+
+proc `$`(a: MyEnum): string =
+  # putting this here pending https://github.com/nim-lang/Nim/issues/13747
+  if a == me2: "me2Modif"
+  else: system.`$`(a)
+
 template fn() = 
   block: # toJson, jsonTo
     type Foo = distinct float
@@ -83,6 +90,12 @@ template fn() =
     doAssert b2.ord == 1 # explains the `1`
     testRoundtrip(a): """[1,2,3]"""
 
+  block: # ToJsonOptions
+    let a = (me1, me2)
+    doAssert $a.toJson() == "[1,2]"
+    doAssert $a.toJson(ToJsonOptions(enumMode: joptEnumSymbol)) == """["me1","me2"]"""
+    doAssert $a.toJson(ToJsonOptions(enumMode: joptEnumString)) == """["me1Alt","me2Modif"]"""
+
   block: # set
     type Foo = enum f1, f2, f3, f4, f5
     type Goo = enum g1 = 10, g2 = 15, g3 = 17, g4