diff --git a/src/services/documentRegistry.ts b/src/services/documentRegistry.ts
index 5720c69be68..4744f25ac7b 100644
--- a/src/services/documentRegistry.ts
+++ b/src/services/documentRegistry.ts
@@ -121,14 +121,6 @@ namespace ts {
         const buckets = createMap<Map<DocumentRegistryEntry>>();
         const getCanonicalFileName = createGetCanonicalFileName(!!useCaseSensitiveFileNames);
 
-        function getBucketForCompilationSettings(key: DocumentRegistryBucketKey, createIfMissing: boolean): Map<DocumentRegistryEntry> {
-            let bucket = buckets.get(key);
-            if (!bucket && createIfMissing) {
-                buckets.set(key, bucket = createMap<DocumentRegistryEntry>());
-            }
-            return bucket!; // TODO: GH#18217
-        }
-
         function reportStats() {
             const bucketInfoArray = arrayFrom(buckets.keys()).filter(name => name && name.charAt(0) === "_").map(name => {
                 const entries = buckets.get(name)!;
@@ -178,7 +170,7 @@ namespace ts {
             acquiring: boolean,
             scriptKind?: ScriptKind): SourceFile {
 
-            const bucket = getBucketForCompilationSettings(key, /*createIfMissing*/ true);
+            const bucket = getOrUpdate<Map<DocumentRegistryEntry>>(buckets, key, createMap);
             let entry = bucket.get(path);
             const scriptTarget = scriptKind === ScriptKind.JSON ? ScriptTarget.JSON : compilationSettings.target || ScriptTarget.ES5;
             if (!entry && externalCache) {
@@ -238,9 +230,7 @@ namespace ts {
         }
 
         function releaseDocumentWithKey(path: Path, key: DocumentRegistryBucketKey): void {
-            const bucket = getBucketForCompilationSettings(key, /*createIfMissing*/ false);
-            Debug.assert(bucket !== undefined);
-
+            const bucket = Debug.assertDefined(buckets.get(key));
             const entry = bucket.get(path)!;
             entry.languageServiceRefCount--;