From 551d76962dfc0ceabf9970a32395d55733ffbe07 Mon Sep 17 00:00:00 2001
From: idkravitz <idkravitz@gmail.com>
Date: Wed, 21 Jun 2023 01:23:38 +1000
Subject: [PATCH] Add filtering of selectors by facets in Diamonds

---
 src/helpers.ts | 14 ++++++++++----
 src/utils.ts   |  9 ++++++++-
 types.ts       |  3 +++
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/src/helpers.ts b/src/helpers.ts
index 28bf18cb..d72a234c 100644
--- a/src/helpers.ts
+++ b/src/helpers.ts
@@ -40,7 +40,7 @@ import {
 } from '../types';
 import {PartialExtension} from './internal/types';
 import {UnknownSignerError} from './errors';
-import {mergeABIs, recode} from './utils';
+import {filterABI, mergeABIs, recode} from './utils';
 import fs from 'fs-extra';
 
 import OpenZeppelinTransparentProxy from '../extendedArtifacts/TransparentUpgradeableProxy.json';
@@ -1958,6 +1958,7 @@ Note that in this case, the contract deployment will not behave the same if depl
     let abi: any[] = diamondArtifact.abi.concat([]);
     const facetCuts: FacetCut[] = [];
     let facetFound: string | undefined;
+    const excludeSelectors: Record<string, string[]> = options.excludeSelectors || {};
     for (const facet of facetsSet) {
       let deterministicFacet: string | boolean = true;
       let facetName;
@@ -2020,7 +2021,12 @@ Note that in this case, the contract deployment will not behave the same if depl
         // reset args for case where facet do not expect any and there was no specific args set on it
         facetArgs = [];
       }
-      abi = mergeABIs([abi, artifact.abi], {
+      let excludeSighashes: Set<string> = new Set();
+      if (facetName in excludeSelectors) {
+        const iface = new Interface(artifact.abi);
+        excludeSighashes = new Set(excludeSelectors[facetName].map(selector => iface.getSighash(selector)));
+      }
+      abi = mergeABIs([abi, filterABI(artifact.abi, excludeSighashes)], {
         check: true,
         skipSupportsInterface: false,
       });
@@ -2045,7 +2051,7 @@ Note that in this case, the contract deployment will not behave the same if depl
         facetAddress = implementation.address;
         const newFacet = {
           facetAddress,
-          functionSelectors: sigsFromABI(implementation.abi),
+          functionSelectors: sigsFromABI(filterABI(implementation.abi, excludeSighashes)),
         };
         facetSnapshot.push(newFacet);
         newSelectors.push(...newFacet.functionSelectors);
@@ -2054,7 +2060,7 @@ Note that in this case, the contract deployment will not behave the same if depl
         facetAddress = oldImpl.address;
         const newFacet = {
           facetAddress,
-          functionSelectors: sigsFromABI(oldImpl.abi),
+          functionSelectors: sigsFromABI(filterABI(oldImpl.abi, excludeSighashes)),
         };
         facetSnapshot.push(newFacet);
         newSelectors.push(...newFacet.functionSelectors);
diff --git a/src/utils.ts b/src/utils.ts
index e1e3be2e..1419fec3 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -6,7 +6,7 @@ import {getAddress, isAddress} from '@ethersproject/address';
 import {Interface, FunctionFragment, Fragment} from '@ethersproject/abi';
 import {Artifact, HardhatRuntimeEnvironment, Network} from 'hardhat/types';
 import {BigNumber} from '@ethersproject/bignumber';
-import {Export, ExtendedArtifact, MultiExport} from '../types';
+import {ABI, Export, ExtendedArtifact, MultiExport} from '../types';
 import {Artifacts} from 'hardhat/internal/artifacts';
 import murmur128 from 'murmur-128';
 import {Transaction} from '@ethersproject/transactions';
@@ -563,6 +563,13 @@ export function getDeployPaths(network: Network): string[] {
   }
 }
 
+export function filterABI(
+  abi: ABI,
+  excludeSighashes: Set<string>,
+): any[] {
+  return abi.filter(fragment => fragment.type !== 'function' || !excludeSighashes.has(Interface.getSighash(Fragment.from(fragment) as FunctionFragment)));
+}
+
 export function mergeABIs(
   abis: any[][],
   options: {check: boolean; skipSupportsInterface: boolean}
diff --git a/types.ts b/types.ts
index aae0311b..d1e73fd9 100644
--- a/types.ts
+++ b/types.ts
@@ -96,6 +96,9 @@ export interface DiamondOptions extends TxOptions {
     methodName: string;
     args: any[];
   };
+  excludeSelectors?: {
+    [facetName: string]: string[]
+  };
   deterministicSalt?: string;
   facetsArgs?: any[];
 }