Allow StoredIntegerElement to use full storage when necessary. (#61726)

* Allow StoredIntegerElement to use full storage when necessary.

This way we will be able to store whole 64-bit number when we store
only a single value, or two 32-bit numbers or 3 16-bit.

Fixes #61723

Author: asl

include/swift/Sema/ConstraintLocator.h

@@ -508,6 +508,26 @@ public: \
 /// A base class for custom path elements that store numeric values.
 template <unsigned NumValues>
 class StoredIntegerElement: public LocatorPathElt {
+  static constexpr unsigned valueWidth() {
+    switch (NumValues) {
+    default:
+      return 16;
+    case 1:
+      return 64;
+    case 2:
+      return 32;
+    }
+  }
+
+  template <unsigned Index = 0,
+            typename = typename std::enable_if<(Index < NumValues)>::type>
+  static uint64_t packValue(unsigned value) {
+    return uint64_t(value) << (valueWidth() * (NumValues - Index - 1));
+  }
+
+  static constexpr uint64_t valueMask =
+    valueWidth() < 64 ? (1ULL << valueWidth()) - 1 : -1ULL;
+
 public:
   template <unsigned NumNumericInputs = NumValues,
             typename = typename std::enable_if<NumNumericInputs == 1>::type>
@@ -519,7 +539,7 @@ class StoredIntegerElement: public LocatorPathElt {
   template <unsigned NumNumericInputs = NumValues,
             typename = typename std::enable_if<NumNumericInputs == 2>::type>
   StoredIntegerElement(ConstraintLocator::PathElementKind kind, unsigned value0, unsigned value1)
-    : LocatorPathElt(kind, (value0 << 16 | value1)) {
+    : LocatorPathElt(kind, packValue<0>(value0) | packValue<1>(value1)) {
     assert(value0 == getValue<0>() && "value0 truncated");
     assert(value1 == getValue<1>() && "value1 truncated");
   }
@@ -528,7 +548,7 @@ class StoredIntegerElement: public LocatorPathElt {
             typename = typename std::enable_if<NumNumericInputs == 3>::type>
   StoredIntegerElement(ConstraintLocator::PathElementKind kind, unsigned value0,
                        unsigned value1, unsigned value2)
-      : LocatorPathElt(kind, uint64_t(value0) << 32 | uint64_t(value1) << 16 | uint64_t(value2)) {
+    : LocatorPathElt(kind, packValue<0>(value0) | packValue<1>(value1) | packValue<2>(value2)) {
     assert(value0 == getValue<0>() && "value0 truncated");
     assert(value1 == getValue<1>() && "value1 truncated");
     assert(value2 == getValue<2>() && "value2 truncated");
@@ -537,13 +557,13 @@ class StoredIntegerElement: public LocatorPathElt {
   /// Retrieve a value associated with the path element.
   template <unsigned Index = 0,
             typename = typename std::enable_if<(Index < NumValues)>::type>
-  unsigned getValue() const {
-    // We pack values into 16 bit components of the storage, with value0
+  uint64_t getValue() const {
+    // We pack values into equally-sized components of the storage, with value0
     // being stored in the upper bits, valueN in the lower bits. Therefore we
     // need to shift out any extra values in the lower bits.
     auto extraValues = NumValues - Index - 1;
-    auto value = getRawStorage() >> (extraValues * 16);
-    return value & 0xFFFF;
+    auto value = getRawStorage() >> (extraValues * valueWidth());
+    return value & valueMask;
   }
 };
 

test/Sema/large_int_array.swift.gyb

@@ -0,0 +1,12 @@
+// RUN: %empty-directory(%t)
+// RUN: %gyb %s -o %t/large_int_array.swift
+// RUN: %target-swift-frontend -typecheck -verify %t/large_int_array.swift
+
+% num_elements = 65537
+
+let v = 0
+let _ : [Int] = [
+% for i in range(num_elements):
+  v,
+% end
+]