From de9ed4e236b0f0db82ce84ee6ff22e356d5dd9dd Mon Sep 17 00:00:00 2001
From: Mike Donnalley <mdonnalley@salesforce.com>
Date: Fri, 11 Oct 2024 10:05:46 -0600
Subject: [PATCH] fix: remove circular references when colorizing json

---
 src/ux/colorize-json.ts       | 38 +++++++++++++++++-
 test/ux/colorize-json.test.ts | 76 ++++++++++++++++++++++++++++++++++-
 2 files changed, 112 insertions(+), 2 deletions(-)

diff --git a/src/ux/colorize-json.ts b/src/ux/colorize-json.ts
index d509475a0..c7171aa31 100644
--- a/src/ux/colorize-json.ts
+++ b/src/ux/colorize-json.ts
@@ -17,6 +17,42 @@ type Options = {
   theme?: Record<string, string> | undefined
 }
 
+export function removeCycles(object: unknown) {
+  // Keep track of seen objects.
+  const seenObjects = new WeakMap<Record<string, unknown>, undefined>()
+
+  const _removeCycles = (obj: unknown) => {
+    // Use object prototype to get around type and null checks
+    if (Object.prototype.toString.call(obj) === '[object Object]') {
+      // We know it is a "Record<string, unknown>" because of the conditional
+      const dictionary = obj as Record<string, unknown>
+
+      // Seen, return undefined to remove.
+      if (seenObjects.has(dictionary)) return
+
+      seenObjects.set(dictionary, undefined)
+
+      for (const key in dictionary) {
+        // Delete the duplicate object if cycle found.
+        if (_removeCycles(dictionary[key]) === undefined) {
+          delete dictionary[key]
+        }
+      }
+    } else if (Array.isArray(obj)) {
+      for (const i in obj) {
+        if (_removeCycles(obj[i]) === undefined) {
+          // We don't want to delete the array, but we can replace the element with null.
+          obj[i] = null
+        }
+      }
+    }
+
+    return obj
+  }
+
+  return _removeCycles(object)
+}
+
 function formatInput(json?: unknown, options?: Options) {
   return options?.pretty
     ? JSON.stringify(typeof json === 'string' ? JSON.parse(json) : json, null, 2)
@@ -26,7 +62,7 @@ function formatInput(json?: unknown, options?: Options) {
 }
 
 export function tokenize(json?: unknown, options?: Options) {
-  let input = formatInput(json, options)
+  let input = formatInput(removeCycles(json), options)
 
   const tokens = []
   let foundToken = false
diff --git a/test/ux/colorize-json.test.ts b/test/ux/colorize-json.test.ts
index a2eb2b1a3..57a89cce9 100644
--- a/test/ux/colorize-json.test.ts
+++ b/test/ux/colorize-json.test.ts
@@ -1,6 +1,6 @@
 import {expect} from 'chai'
 
-import {tokenize} from '../../src/ux/colorize-json'
+import {removeCycles, tokenize} from '../../src/ux/colorize-json'
 
 describe('colorizeJson', () => {
   it('tokenizes a basic JSON object', () => {
@@ -124,4 +124,78 @@ describe('colorizeJson', () => {
 
     expect(result).to.deep.equal([])
   })
+
+  it('removes circular references from json', () => {
+    const obj = {
+      foo: 'bar',
+      baz: {
+        qux: 'quux',
+      },
+    }
+    // @ts-expect-error
+    obj.circular = obj
+
+    const result = tokenize(obj)
+    expect(result).to.deep.equal([
+      {type: 'brace', value: '{'},
+      {type: 'key', value: '"foo"'},
+      {type: 'colon', value: ':'},
+      {type: 'string', value: '"bar"'},
+      {type: 'comma', value: ','},
+      {type: 'key', value: '"baz"'},
+      {type: 'colon', value: ':'},
+      {type: 'brace', value: '{'},
+      {type: 'key', value: '"qux"'},
+      {type: 'colon', value: ':'},
+      {type: 'string', value: '"quux"'},
+      {type: 'brace', value: '}'},
+      {type: 'brace', value: '}'},
+    ])
+  })
+})
+
+describe('removeCycles', () => {
+  it('removes circular references from objects', () => {
+    const obj = {
+      foo: 'bar',
+      baz: {
+        qux: 'quux',
+      },
+    }
+    // @ts-expect-error
+    obj.circular = obj
+
+    const result = removeCycles(obj)
+    expect(result).to.deep.equal({
+      foo: 'bar',
+      baz: {
+        qux: 'quux',
+      },
+    })
+  })
+
+  it('removes circular references from objects in array', () => {
+    const obj = {
+      foo: 'bar',
+      baz: {
+        qux: 'quux',
+      },
+    }
+    // @ts-expect-error
+    obj.circular = obj
+    const arr = [{foo: 'bar'}, obj]
+
+    const result = removeCycles(arr)
+    expect(result).to.deep.equal([
+      {
+        foo: 'bar',
+      },
+      {
+        baz: {
+          qux: 'quux',
+        },
+        foo: 'bar',
+      },
+    ])
+  })
 })