refactor: clean up some of the light account types (#828)

alchemyplatform · Jul 30, 2024 · 08048fc · 08048fc
1 parent bfcf4f7
commit 08048fc
Show file tree

Hide file tree

Showing 13 changed files with 233 additions and 615 deletions.
diff --git a/account-kit/core/src/actions/createAccount.ts b/account-kit/core/src/actions/createAccount.ts
@@ -1,11 +1,11 @@
+import { type SmartAccountSigner } from "@aa-sdk/core";
 import {
   createLightAccount,
   createMultiOwnerModularAccount,
   type CreateLightAccountParams,
   type CreateMultiOwnerModularAccountParams,
-  type GetLightAccountVersion,
+  type LightAccountVersion,
 } from "@account-kit/smart-contracts";
-import { type SmartAccountSigner } from "@aa-sdk/core";
 import { custom, type Transport } from "viem";
 import { ClientOnlyPropertyError } from "../errors.js";
 import type {
@@ -23,8 +23,7 @@ export type AccountConfig<TAccount extends SupportedAccountTypes> =
         CreateLightAccountParams<
           Transport,
           SmartAccountSigner,
-          // we only support LightAccount version v1
-          Exclude<GetLightAccountVersion<"LightAccount">, "v2.0.0">
+          LightAccountVersion<"LightAccount">
         >,
         "signer" | "transport" | "chain"
       >
@@ -94,8 +93,11 @@ export async function createAccount<TAccount extends SupportedAccountTypes>(
         });
       case "MultiOwnerModularAccount":
         return createMultiOwnerModularAccount({
-          ...params,
-          ...cachedConfig,
+          ...(params as AccountConfig<"MultiOwnerModularAccount">),
+          ...(cachedConfig as Omit<
+            CreateMultiOwnerModularAccountParams,
+            "transport" | "chain" | "signer"
+          >),
           signer,
           transport: (opts) => transport({ ...opts, retryCount: 0 }),
           chain,

diff --git a/account-kit/smart-contracts/src/index.ts b/account-kit/smart-contracts/src/index.ts
@@ -24,7 +24,7 @@ export {
   LightAccountUnsupported1271Impls,
   defaultLightAccountVersion,
   getDefaultLightAccountFactoryAddress,
-  getLightAccountVersionDef,
+  getLightAccountVersionForAccount,
 } from "./light-account/utils.js";
 
 //multi-owner-light-account exports

diff --git a/account-kit/smart-contracts/src/light-account/accounts/account.test.ts b/account-kit/smart-contracts/src/light-account/accounts/account.test.ts
@@ -13,7 +13,7 @@ const chain = polygonMumbai;
 
 const versions = Object.keys(
   AccountVersionRegistry.LightAccount
-) as LightAccountVersion[];
+) as LightAccountVersion<"LightAccount">[];
 
 describe("Light Account Tests", () => {
   const dummyMnemonic =
@@ -170,7 +170,7 @@ const givenConnectedProvider = ({
 }: {
   signer: SmartAccountSigner;
   chain: Chain;
-  version: LightAccountVersion;
+  version: LightAccountVersion<"LightAccount">;
 }) =>
   createLightAccountClient({
     account: {
@@ -181,8 +181,12 @@ const givenConnectedProvider = ({
     transport: custom({
       request: async ({ method }) => {
         if (method === "eth_getStorageAt") {
-          return AccountVersionRegistry.LightAccount[version].address[chain.id]
-            .impl;
+          return (
+            AccountVersionRegistry.LightAccount[version].addresses.overrides?.[
+              chain.id
+            ].impl ??
+            AccountVersionRegistry.LightAccount[version].addresses.default.impl
+          );
         }
         return;
       },

diff --git a/account-kit/smart-contracts/src/light-account/accounts/account.ts b/account-kit/smart-contracts/src/light-account/accounts/account.ts
@@ -18,8 +18,8 @@ import { LightAccountAbi_v2 } from "../abis/LightAccountAbi_v2.js";
 import { LightAccountFactoryAbi_v1 } from "../abis/LightAccountFactoryAbi_v1.js";
 import { LightAccountFactoryAbi_v2 } from "../abis/LightAccountFactoryAbi_v2.js";
 import type {
-  GetEntryPointForLightAccountVersion,
-  GetLightAccountVersion,
+  LightAccountEntryPointVersion,
+  LightAccountVersion,
 } from "../types.js";
 import {
   AccountVersionRegistry,
@@ -35,17 +35,8 @@ import {
 
 export type LightAccount<
   TSigner extends SmartAccountSigner = SmartAccountSigner,
-  TLightAccountVersion extends GetLightAccountVersion<"LightAccount"> = GetLightAccountVersion<"LightAccount">,
-  TEntryPointVersion extends GetEntryPointForLightAccountVersion<
-    "LightAccount",
-    TLightAccountVersion
-  > = GetEntryPointForLightAccountVersion<"LightAccount", TLightAccountVersion>
-> = LightAccountBase<
-  TSigner,
-  "LightAccount",
-  TLightAccountVersion,
-  TEntryPointVersion
-> & {
+  TLightAccountVersion extends LightAccountVersion<"LightAccount"> = LightAccountVersion<"LightAccount">
+> = LightAccountBase<TSigner, "LightAccount", TLightAccountVersion> & {
   encodeTransferOwnership: (newOwner: Address) => Hex;
   getOwnerAddress: () => Promise<Address>;
 };
@@ -54,55 +45,41 @@ export type LightAccount<
 export type CreateLightAccountParams<
   TTransport extends Transport = Transport,
   TSigner extends SmartAccountSigner = SmartAccountSigner,
-  TLightAccountVersion extends GetLightAccountVersion<"LightAccount"> = GetLightAccountVersion<"LightAccount">,
-  TEntryPointVersion extends GetEntryPointForLightAccountVersion<
-    "LightAccount",
-    TLightAccountVersion
-  > = GetEntryPointForLightAccountVersion<"LightAccount", TLightAccountVersion>
+  TLightAccountVersion extends LightAccountVersion<"LightAccount"> = LightAccountVersion<"LightAccount">
 > = Omit<
   CreateLightAccountBaseParams<
-    TTransport,
-    TSigner,
     "LightAccount",
     TLightAccountVersion,
-    TEntryPointVersion
+    TTransport,
+    TSigner
   >,
-  "getAccountInitCode" | "entryPoint" | "version" | "abi" | "accountAddress"
+  | "getAccountInitCode"
+  | "entryPoint"
+  | "version"
+  | "abi"
+  | "accountAddress"
+  | "type"
 > & {
   salt?: bigint;
   initCode?: Hex;
   accountAddress?: Address;
   factoryAddress?: Address;
   version?: TLightAccountVersion;
-  entryPoint?: EntryPointDef<TEntryPointVersion, Chain>;
+  entryPoint?: EntryPointDef<
+    LightAccountEntryPointVersion<"LightAccount", TLightAccountVersion>,
+    Chain
+  >;
 };
 // [!endregion CreateLightAccountParams]
 
 export async function createLightAccount<
   TTransport extends Transport = Transport,
   TSigner extends SmartAccountSigner = SmartAccountSigner,
-  TLightAccountVersion extends GetLightAccountVersion<"LightAccount"> = "v1.1.0"
+  TLightAccountVersion extends LightAccountVersion<"LightAccount"> = "v2.0.0"
 >(
   config: CreateLightAccountParams<TTransport, TSigner, TLightAccountVersion>
 ): Promise<LightAccount<TSigner, TLightAccountVersion>>;
 
-export async function createLightAccount<
-  TTransport extends Transport = Transport,
-  TSigner extends SmartAccountSigner = SmartAccountSigner,
-  TLightAccountVersion extends GetLightAccountVersion<"LightAccount"> = GetLightAccountVersion<"LightAccount">,
-  TEntryPointVersion extends GetEntryPointForLightAccountVersion<
-    "LightAccount",
-    TLightAccountVersion
-  > = GetEntryPointForLightAccountVersion<"LightAccount", TLightAccountVersion>
->(
-  config: CreateLightAccountParams<
-    TTransport,
-    TSigner,
-    TLightAccountVersion,
-    TEntryPointVersion
-  >
-): Promise<LightAccount<TSigner, TLightAccountVersion, TEntryPointVersion>>;
-
 /**
  * Creates a light account based on the provided parameters such as transport, chain, signer, init code, and more. Ensures that an account is configured and returned with various capabilities, such as transferring ownership and retrieving the owner's address.
  *
@@ -128,7 +105,7 @@ export async function createLightAccount({
   chain,
   signer,
   initCode,
-  version = defaultLightAccountVersion("LightAccount"),
+  version = defaultLightAccountVersion(),
   entryPoint = getEntryPoint(chain, {
     version: AccountVersionRegistry["LightAccount"][version]
       .entryPointVersion as any,
@@ -176,15 +153,17 @@ export async function createLightAccount({
   });
 
   const account = await createLightAccountBase<
+    "LightAccount",
+    LightAccountVersion<"LightAccount">,
     Transport,
-    SmartAccountSigner,
-    "LightAccount"
+    SmartAccountSigner
   >({
     transport,
     chain,
     signer,
     abi: accountAbi,
-    version: AccountVersionRegistry["LightAccount"][version],
+    type: "LightAccount",
+    version,
     entryPoint,
     accountAddress: address,
     getAccountInitCode,