diff --git a/clang-tools-extra/clangd/benchmarks/CompletionModel/DecisionForestBenchmark.cpp b/clang-tools-extra/clangd/benchmarks/CompletionModel/DecisionForestBenchmark.cpp
--- a/clang-tools-extra/clangd/benchmarks/CompletionModel/DecisionForestBenchmark.cpp
+++ b/clang-tools-extra/clangd/benchmarks/CompletionModel/DecisionForestBenchmark.cpp
@@ -51,7 +51,7 @@
                       : RandInt(20));
     E.setSemaSaysInScope(FlipCoin(0.5));      // Boolean.
     E.setScope(RandInt(4));                   // 4 Scopes.
-    E.setContextKind(RandInt(32));            // 32 Context kinds.
+    E.setContextKind(RandInt(36));            // 36 Context kinds.
     E.setIsInstanceMember(FlipCoin(0.5));     // Boolean.
     E.setHadContextType(FlipCoin(0.6));       // Boolean.
     E.setHadSymbolType(FlipCoin(0.6));        // Boolean.
diff --git a/clang-tools-extra/clangd/quality/CompletionModelCodegen.py b/clang-tools-extra/clangd/quality/CompletionModelCodegen.py
--- a/clang-tools-extra/clangd/quality/CompletionModelCodegen.py
+++ b/clang-tools-extra/clangd/quality/CompletionModelCodegen.py
@@ -131,18 +131,19 @@
                     feature, feature))
         elif f["kind"] == "ENUM":
             setters.append(
-                "void set%s(unsigned V) { %s = 1 << V; }" % (feature, feature))
+                "void set%s(unsigned V) { %s = 1LL << V; }" % (feature, feature))
         else:
             raise ValueError("Unhandled feature type.", f["kind"])
 
     # Class members represent all the features of the Example.
     class_members = [
-        "uint32_t %s = 0;" % f['name']
+        "uint%d_t %s = 0;"
+        % (64 if f["kind"] == "ENUM" else 32, f['name'])
         for f in features_json
     ]
     getters = [
-        "LLVM_ATTRIBUTE_ALWAYS_INLINE uint32_t get%s() const { return %s; }"
-        % (f['name'], f['name'])
+        "LLVM_ATTRIBUTE_ALWAYS_INLINE uint%d_t get%s() const { return %s; }"
+        % (64 if f["kind"] == "ENUM" else 32, f['name'], f['name'])
         for f in features_json
     ]
     nline = "\n  "
@@ -245,7 +246,7 @@
 
 %s
 
-#define BIT(X) (1 << X)
+#define BIT(X) (1LL << X)
 
 %s