From a7edda4b1acbbbb2e27e892715b6882b0be8e945 Mon Sep 17 00:00:00 2001
From: 0xCourtney <23440493+0xCourtney@users.noreply.github.com>
Date: Thu, 10 Nov 2022 15:15:39 -0500
Subject: [PATCH] diamond utils draft

---
 lib/diamonds.ts  | 327 +++++++++++++++++++++++++++++++++++++++++++++++
 lib/index.ts     |   1 +
 lib/package.json |   1 +
 3 files changed, 329 insertions(+)
 create mode 100644 lib/diamonds.ts

diff --git a/lib/diamonds.ts b/lib/diamonds.ts
new file mode 100644
index 00000000..b8b0a4a1
--- /dev/null
+++ b/lib/diamonds.ts
@@ -0,0 +1,327 @@
+import {
+  IDiamondReadable,
+  IDiamondWritable,
+} from '@solidstate/typechain-types';
+import { Contract, constants } from 'ethers';
+
+const Table = require('cli-table3');
+
+type Diamond = IDiamondReadable & IDiamondWritable;
+
+export interface Facet {
+  target: string;
+  selectors: string[];
+}
+
+enum FacetCutAction {
+  Add,
+  Replace,
+  Remove,
+}
+
+export interface FacetCut extends Facet {
+  action: FacetCutAction;
+}
+
+export function getSignatures(contract: Contract): string[] {
+  return Object.keys(contract.interface.functions);
+}
+
+export function getSelectors(contract: Contract): string[] {
+  const signatures = getSignatures(contract);
+  return signatures.reduce((acc: string[], val: string) => {
+    acc.push(contract.interface.getSighash(val));
+    return acc;
+  }, []);
+}
+
+export function getFacets(contracts: Contract[]): Facet[] {
+  return contracts.map((contract) => {
+    return {
+      target: contract.address,
+      selectors: getSelectors(contract),
+    };
+  });
+}
+
+export function selectorExistsInFacets(
+  selector: string,
+  facets: Facet[],
+): boolean {
+  for (const facet of facets) {
+    if (facet.selectors.includes(selector)) return true;
+  }
+
+  return false;
+}
+
+// adds unregistered selectors
+export async function addUnregisteredSelectors(
+  diamond: Diamond,
+  contracts: Contract[],
+  exclude: string[] = [],
+): Promise<FacetCut[]> {
+  const diamondFacets: Facet[] = await diamond.facets();
+  const facets = getFacets(contracts);
+  let facetCuts: FacetCut[] = [];
+
+  // if facet selector is unregistered then it should be added to the diamond.
+  for (const facet of facets) {
+    for (const selector of facet.selectors) {
+      const target = facet.target;
+
+      if (
+        target !== diamond.address &&
+        selector.length > 0 &&
+        !selectorExistsInFacets(selector, diamondFacets) &&
+        !exclude.includes(selector)
+      ) {
+        facetCuts.push(
+          printFacetCuts(facet.target, [selector], FacetCutAction.Add),
+        );
+      }
+    }
+  }
+
+  return groupFacetCuts(facetCuts);
+}
+
+// replace registered selectors
+export async function replaceRegisteredSelectors(
+  diamond: Diamond,
+  contracts: Contract[],
+  exclude: string[] = [],
+): Promise<FacetCut[]> {
+  const diamondFacets: Facet[] = await diamond.facets();
+  const facets = getFacets(contracts);
+  let facetCuts: FacetCut[] = [];
+
+  // if a facet selector is registered with a different target address, the target will be replaced
+  for (const facet of facets) {
+    for (const selector of facet.selectors) {
+      const target = facet.target;
+      const oldTarget = await diamond.facetAddress(selector);
+
+      if (
+        target != oldTarget &&
+        target != constants.AddressZero &&
+        target != diamond.address &&
+        selector.length > 0 &&
+        selectorExistsInFacets(selector, diamondFacets) &&
+        !exclude.includes(selector)
+      ) {
+        facetCuts.push(
+          printFacetCuts(target, [selector], FacetCutAction.Replace),
+        );
+      }
+    }
+  }
+
+  return groupFacetCuts(facetCuts);
+}
+
+// removes registered selectors
+export async function removeRegisteredSelectors(
+  diamond: Diamond,
+  contracts: Contract[],
+  exclude: string[] = [],
+): Promise<FacetCut[]> {
+  const diamondFacets: Facet[] = await diamond.facets();
+  const facets = getFacets(contracts);
+  let facetCuts: FacetCut[] = [];
+
+  // if a registered selector is not found in the facets then it should be removed from the diamond
+  for (const diamondFacet of diamondFacets) {
+    for (const selector of diamondFacet.selectors) {
+      const target = diamondFacet.target;
+
+      if (
+        target != constants.AddressZero &&
+        target != diamond.address &&
+        selector.length > 0 &&
+        !selectorExistsInFacets(selector, facets) &&
+        !exclude.includes(selector)
+      ) {
+        facetCuts.push(
+          printFacetCuts(
+            constants.AddressZero,
+            [selector],
+            FacetCutAction.Remove,
+          ),
+        );
+      }
+    }
+  }
+
+  return groupFacetCuts(facetCuts);
+}
+
+export async function diamondCut(
+  diamond: Diamond,
+  facetCut: FacetCut[],
+  target: string = constants.AddressZero,
+  data: string = '0x',
+) {
+  (await diamond.diamondCut(facetCut, target, data)).wait(1);
+}
+
+// groups facet cuts by target address and action type
+export function groupFacetCuts(facetCuts: FacetCut[]): FacetCut[] {
+  const cuts = facetCuts.reduce((acc: FacetCut[], facetCut: FacetCut) => {
+    if (acc.length == 0) acc.push(facetCut);
+
+    let exists = false;
+
+    acc.forEach((_, i) => {
+      if (
+        acc[i].action == facetCut.action &&
+        acc[i].target == facetCut.target
+      ) {
+        acc[i].selectors.push(...facetCut.selectors);
+        // removes duplicates, if there are any
+        acc[i].selectors = [...new Set(acc[i].selectors)];
+        exists = true;
+      }
+    });
+
+    // push facet cut if it does not already exist
+    if (!exists) acc.push(facetCut);
+
+    return acc;
+  }, []);
+
+  let cache: any = {};
+
+  // checks if selector is used multiple times, emits warning
+  cuts.forEach((cut) => {
+    cut.selectors.forEach((selector: string) => {
+      if (cache[selector]) {
+        console.log(
+          `WARNING: selector: ${selector}, target: ${cut.target} is defined in multiple cuts`,
+        );
+      } else {
+        cache[selector] = true;
+      }
+    });
+  });
+
+  return cuts;
+}
+
+export function printFacetCuts(
+  target: string,
+  selectors: string[],
+  action: number = 0,
+): FacetCut {
+  return {
+    target: target,
+    action: action,
+    selectors: selectors,
+  };
+}
+
+// generates table of diamond and facet selectors
+export async function printDiamond(diamond: Diamond, contracts: Contract[]) {
+  const padding = 2;
+
+  const table = new Table({
+    style: {
+      head: [],
+      border: [],
+      'padding-left': padding,
+      'padding-right': padding,
+    },
+    chars: {
+      mid: '·',
+      'top-mid': '|',
+      'left-mid': ' ·',
+      'mid-mid': '|',
+      'right-mid': '·',
+      left: ' |',
+      'top-left': ' ·',
+      'top-right': '·',
+      'bottom-left': ' ·',
+      'bottom-right': '·',
+      middle: '·',
+      top: '-',
+      bottom: '-',
+      'bottom-mid': '|',
+    },
+  });
+
+  table.push([
+    {
+      hAlign: 'center',
+      content: `Target`,
+    },
+    {
+      hAlign: 'center',
+      content: `Signature`,
+    },
+    {
+      hAlign: 'center',
+      content: `Selector`,
+    },
+    {
+      hAlign: 'center',
+      content: `Registered`,
+    },
+  ]);
+
+  let diamondTable = [];
+  const signatures = await getSignatures(diamond);
+
+  for (const signature of signatures) {
+    diamondTable.push({
+      target: diamond.address,
+      signature: signature,
+      selector: diamond.interface.getSighash(signature),
+      registered: true,
+    });
+  }
+
+  for (const contract of contracts) {
+    const signatures = await getSignatures(contract);
+    for (const signature of signatures) {
+      diamondTable.push({
+        target: contract.address,
+        signature: signature,
+        selector: contract.interface.getSighash(signature),
+        registered: false,
+      });
+    }
+  }
+
+  const diamondFacets: Facet[] = await diamond.facets();
+
+  for (const facet of diamondFacets) {
+    const target = facet.target;
+
+    for (const selector of facet.selectors) {
+      for (const row of diamondTable) {
+        if (row.target == target && row.selector == selector) {
+          row.registered = true;
+        }
+      }
+    }
+  }
+
+  for (const row of diamondTable) {
+    table.push([
+      {
+        content: row.target,
+      },
+      {
+        content: row.signature,
+      },
+      {
+        content: row.selector,
+      },
+      {
+        content: row.registered,
+      },
+    ]);
+  }
+
+  console.log(table.toString());
+}
diff --git a/lib/index.ts b/lib/index.ts
index e6459ee4..3da4ce2d 100644
--- a/lib/index.ts
+++ b/lib/index.ts
@@ -2,3 +2,4 @@ export * from './bn_conversion';
 export * from './erc20_permit';
 export * from './mocha_describe_filter';
 export * from './sign_data';
+export * from './diamonds';
diff --git a/lib/package.json b/lib/package.json
index 0faa17a2..89b4745a 100644
--- a/lib/package.json
+++ b/lib/package.json
@@ -25,6 +25,7 @@
     "tsc-clean": "tsc --build --clean tsconfig.json"
   },
   "dependencies": {
+    "cli-table3": "^0.6.3",
     "eth-permit": "^0.1.10"
   },
   "files": [