Consider range type of runtime discrim [feature] (#11432)

3 files changed, 99 insertions, 17 deletions

diff --git a/compiler/semobjconstr.nim b/compiler/semobjconstr.nim
index 719b5e18d..9f342e3b9 100644
--- a/compiler/semobjconstr.nim
+++ b/compiler/semobjconstr.nim
@@ -113,6 +113,13 @@ proc branchVals(c: PContext, caseNode: PNode, caseIdx: int,
     for i in 1 .. caseNode.len-2:
       processBranchVals(caseNode[i], excl)
 
+proc rangeTypVals(rangeTyp: PType): IntSet =
+  assert rangeTyp.kind == tyRange
+  let (a, b) = (rangeTyp.n.sons[0].intVal, rangeTyp.n.sons[1].intVal)
+  result = initIntSet()
+  for it in a .. b:
+    result.incl(it.int)
+
 proc formatUnsafeBranchVals(t: PType, diffVals: IntSet): string =
   if diffVals.len <= 32:
     var strs: seq[string]
@@ -228,9 +235,9 @@ proc semConstructFields(c: PContext, recNode: PNode,
       template badDiscriminatorError =
         let fields = fieldsPresentInBranch(selectedBranch)
         localError(c.config, initExpr.info,
-          ("you must provide a compile-time value for the discriminator '$1' " &
-          "in order to prove that it's safe to initialize $2.") %
-          [discriminator.sym.name.s, fields])
+          ("cannot prove that it's safe to initialize $1 with " &
+          "the runtime value for the discriminator '$2' ") %
+          [fields, discriminator.sym.name.s])
         mergeInitStatus(result, initNone)
 
       template wrongBranchError(i) =
@@ -241,30 +248,45 @@ proc semConstructFields(c: PContext, recNode: PNode,
           "are in conflict with this value.",
           [discriminator.sym.name.s, discriminatorVal.renderTree, fields])
 
+      template valuesInConflictError(valsDiff) =
+        localError(c.config, discriminatorVal.info, ("possible values " &
+          "$2are in conflict with discriminator values for " &
+          "selected object branch $1.") % [$selectedBranch,
+          formatUnsafeBranchVals(recNode.sons[0].typ, valsDiff)])
+
       let branchNode = recNode[selectedBranch]
       let flags = flags*{efAllowDestructor} + {efPreferStatic,
                                                efPreferNilResult}
       var discriminatorVal = semConstrField(c, flags,
                                             discriminator.sym, initExpr)
-
       if discriminatorVal != nil:
         discriminatorVal = discriminatorVal.skipHidden
+        if discriminatorVal.kind notin nkLiterals and (
+            not isOrdinalType(discriminatorVal.typ, true) or
+            lengthOrd(c.config, discriminatorVal.typ) > MaxSetElements or
+            lengthOrd(c.config, recNode.sons[0].typ) > MaxSetElements):
+          localError(c.config, discriminatorVal.info,
+            "branch initialization with a runtime discriminator only " &
+            "supports ordinal types with 2^16 elements or less.")
+
       if discriminatorVal == nil:
         badDiscriminatorError()
       elif discriminatorVal.kind == nkSym:
         let (ctorCase, ctorIdx) = findUsefulCaseContext(c, discriminatorVal)
         if ctorCase == nil:
-          badDiscriminatorError()
+          if discriminatorVal.typ.kind == tyRange:
+            let rangeVals = rangeTypVals(discriminatorVal.typ)
+            let recBranchVals = branchVals(c, recNode, selectedBranch, false)
+            let diff = rangeVals - recBranchVals
+            if diff.len != 0:
+              valuesInConflictError(diff)
+          else:
+            badDiscriminatorError()
         elif discriminatorVal.sym.kind notin {skLet, skParam} or
             discriminatorVal.sym.typ.kind == tyVar:
           localError(c.config, discriminatorVal.info,
             "runtime discriminator must be immutable if branch fields are " &
             "initialized, a 'let' binding is required.")
-        elif not isOrdinalType(discriminatorVal.sym.typ, true) or
-            lengthOrd(c.config, discriminatorVal.sym.typ) > MaxSetElements:
-          localError(c.config, discriminatorVal.info,
-            "branch initialization with a runtime discriminator only " &
-            "supports ordinal types with 2^16 elements or less.")
         elif ctorCase[ctorIdx].kind == nkElifBranch:
           localError(c.config, discriminatorVal.info, "branch initialization " &
             "with a runtime discriminator is not supported inside of an " &
@@ -275,20 +297,27 @@ proc semConstructFields(c: PContext, recNode: PNode,
             recBranchVals = branchVals(c, recNode, selectedBranch, false)
             branchValsDiff = ctorBranchVals - recBranchVals
           if branchValsDiff.len != 0:
-            localError(c.config, discriminatorVal.info, ("possible values " &
-              "$2are in conflict with discriminator values for " &
-              "selected object branch $1.") % [$selectedBranch,
-              formatUnsafeBranchVals(recNode.sons[0].typ, branchValsDiff)])
+            valuesInConflictError(branchValsDiff)
       else:
+        var failedBranch = -1
         if branchNode.kind != nkElse:
           if not branchNode.caseBranchMatchesExpr(discriminatorVal):
-            wrongBranchError(selectedBranch)
+            failedBranch = selectedBranch
         else:
           # With an else clause, check that all other branches don't match:
           for i in 1 .. (recNode.len - 2):
             if recNode[i].caseBranchMatchesExpr(discriminatorVal):
-              wrongBranchError(i)
+              failedBranch = i
               break
+        if failedBranch != -1:
+          if discriminatorVal.typ.kind == tyRange:
+            let rangeVals = rangeTypVals(discriminatorVal.typ)
+            let recBranchVals = branchVals(c, recNode, selectedBranch, false)
+            let diff = rangeVals - recBranchVals
+            if diff.len != 0:
+              valuesInConflictError(diff)
+          else:
+            wrongBranchError(failedBranch)
 
       # When a branch is selected with a partial match, some of the fields
       # that were not initialized may be mandatory. We must check for this:
diff --git a/doc/manual.rst b/doc/manual.rst
index 0e4eb08a8..e1fb2fa59 100644
--- a/doc/manual.rst
+++ b/doc/manual.rst
@@ -1600,7 +1600,9 @@ statement. If possible values of the discriminator variable in a
 ``case`` statement branch are a subset of discriminator values for the selected
 object branch, the initialization is considered valid. This analysis only works
 for immutable discriminators of an ordinal type and disregards ``elif``
-branches.
+branches. For discriminator values with a ``range`` type, the compiler
+checks if the entire range of possible values for the discriminator value is
+valid for the choosen object branch.
 
 A small example:
 
@@ -1619,6 +1621,10 @@ A small example:
   else:
     echo "ignoring: ", unknownKind
 
+  # also valid, since unknownKindBounded can only contain the values nkAdd or nkSub
+  let unknownKindBounded = range[nkAdd..nkSub](unknownKind)
+  z = Node(kind: unknownKindBounded, leftOp: Node(), rightOp: Node())
+
 Set type
 --------
 
diff --git a/tests/objvariant/trt_discrim.nim b/tests/objvariant/trt_discrim.nim
index 303c0fa55..30ba3d087 100644
--- a/tests/objvariant/trt_discrim.nim
+++ b/tests/objvariant/trt_discrim.nim
@@ -148,3 +148,50 @@ reject:
     case kind:
     of k1: result = KindObj(kind: kind, i32: 1)
     else: discard
+    
+type
+  Kind3 = enum
+    A, B, C, E
+
+  OkRange = range[B..C]
+  NotOkRange = range[B..E]
+
+  CaseObject = object
+    case kind: Kind3
+    of B, C:
+      field: int
+    else: discard
+
+accept:
+  let rtDiscriminator: OkRange = B
+  discard CaseObject(kind: rtDiscriminator, field: 1)
+
+accept:
+  let rtDiscriminator = B
+  discard CaseObject(kind: OkRange(rtDiscriminator), field: 1)
+
+accept:
+  const rtDiscriminator: NotOkRange = B
+  discard CaseObject(kind: rtDiscriminator, field: 1)
+
+accept:
+  discard CaseObject(kind: NotOkRange(B), field: 1)
+
+reject:
+  let rtDiscriminator: NotOkRange = B
+  discard CaseObject(kind: rtDiscriminator, field: 1)
+
+reject:
+  let rtDiscriminator = B
+  discard CaseObject(kind: NotOkRange(rtDiscriminator), field: 1)
+
+reject:
+  type Obj = object
+    case x: int
+    of 0 .. 1000:
+      field: int
+    else:
+      discard
+
+  let x: range[0..15] = 1
+  let o = Obj(x: x, field: 1)
\ No newline at end of file