feat: implement message v0 compilation (#27524)

* feat: add PublicKey.unique method for tests

* feat: add MessageAccountKeys class

* feat: add CompiledKeys class for message compilation

* feat: implement message compilation using CompiledKeys

web3.js/src/message/account-keys.ts

@@ -0,0 +1,79 @@
+import {LoadedAddresses} from '../connection';
+import {PublicKey} from '../publickey';
+import {TransactionInstruction} from '../transaction';
+import {MessageCompiledInstruction} from './index';
+
+export type AccountKeysFromLookups = LoadedAddresses;
+
+export class MessageAccountKeys {
+  staticAccountKeys: Array<PublicKey>;
+  accountKeysFromLookups?: AccountKeysFromLookups;
+
+  constructor(
+    staticAccountKeys: Array<PublicKey>,
+    accountKeysFromLookups?: AccountKeysFromLookups,
+  ) {
+    this.staticAccountKeys = staticAccountKeys;
+    this.accountKeysFromLookups = accountKeysFromLookups;
+  }
+
+  keySegments(): Array<Array<PublicKey>> {
+    const keySegments = [this.staticAccountKeys];
+    if (this.accountKeysFromLookups) {
+      keySegments.push(this.accountKeysFromLookups.writable);
+      keySegments.push(this.accountKeysFromLookups.readonly);
+    }
+    return keySegments;
+  }
+
+  get(index: number): PublicKey | undefined {
+    for (const keySegment of this.keySegments()) {
+      if (index < keySegment.length) {
+        return keySegment[index];
+      } else {
+        index -= keySegment.length;
+      }
+    }
+    return;
+  }
+
+  get length(): number {
+    return this.keySegments().flat().length;
+  }
+
+  compileInstructions(
+    instructions: Array<TransactionInstruction>,
+  ): Array<MessageCompiledInstruction> {
+    // Bail early if any account indexes would overflow a u8
+    const U8_MAX = 255;
+    if (this.length > U8_MAX + 1) {
+      throw new Error('Account index overflow encountered during compilation');
+    }
+
+    const keyIndexMap = new Map();
+    this.keySegments()
+      .flat()
+      .forEach((key, index) => {
+        keyIndexMap.set(key.toBase58(), index);
+      });
+
+    const findKeyIndex = (key: PublicKey) => {
+      const keyIndex = keyIndexMap.get(key.toBase58());
+      if (keyIndex === undefined)
+        throw new Error(
+          'Encountered an unknown instruction account key during compilation',
+        );
+      return keyIndex;
+    };
+
+    return instructions.map((instruction): MessageCompiledInstruction => {
+      return {
+        programIdIndex: findKeyIndex(instruction.programId),
+        accountKeyIndexes: instruction.keys.map(meta =>
+          findKeyIndex(meta.pubkey),
+        ),
+        data: instruction.data,
+      };
+    });
+  }
+}

web3.js/src/message/compiled-keys.ts

@@ -0,0 +1,165 @@
+import {MessageHeader, MessageAddressTableLookup} from './index';
+import {AccountKeysFromLookups} from './account-keys';
+import {AddressLookupTableAccount} from '../programs';
+import {TransactionInstruction} from '../transaction';
+import assert from '../utils/assert';
+import {PublicKey} from '../publickey';
+
+export type CompiledKeyMeta = {
+  isSigner: boolean;
+  isWritable: boolean;
+  isInvoked: boolean;
+};
+
+type KeyMetaMap = Map<string, CompiledKeyMeta>;
+
+export class CompiledKeys {
+  payer: PublicKey;
+  keyMetaMap: KeyMetaMap;
+
+  constructor(payer: PublicKey, keyMetaMap: KeyMetaMap) {
+    this.payer = payer;
+    this.keyMetaMap = keyMetaMap;
+  }
+
+  static compile(
+    instructions: Array<TransactionInstruction>,
+    payer: PublicKey,
+  ): CompiledKeys {
+    const keyMetaMap: KeyMetaMap = new Map();
+    const getOrInsertDefault = (pubkey: PublicKey): CompiledKeyMeta => {
+      const address = pubkey.toBase58();
+      let keyMeta = keyMetaMap.get(address);
+      if (keyMeta === undefined) {
+        keyMeta = {
+          isSigner: false,
+          isWritable: false,
+          isInvoked: false,
+        };
+        keyMetaMap.set(address, keyMeta);
+      }
+      return keyMeta;
+    };
+
+    const payerKeyMeta = getOrInsertDefault(payer);
+    payerKeyMeta.isSigner = true;
+    payerKeyMeta.isWritable = true;
+
+    for (const ix of instructions) {
+      getOrInsertDefault(ix.programId).isInvoked = true;
+      for (const accountMeta of ix.keys) {
+        const keyMeta = getOrInsertDefault(accountMeta.pubkey);
+        keyMeta.isSigner ||= accountMeta.isSigner;
+        keyMeta.isWritable ||= accountMeta.isWritable;
+      }
+    }
+
+    return new CompiledKeys(payer, keyMetaMap);
+  }
+
+  getMessageComponents(): [MessageHeader, Array<PublicKey>] {
+    const mapEntries = [...this.keyMetaMap.entries()];
+    assert(mapEntries.length <= 256, 'Max static account keys length exceeded');
+
+    const writableSigners = mapEntries.filter(
+      ([, meta]) => meta.isSigner && meta.isWritable,
+    );
+    const readonlySigners = mapEntries.filter(
+      ([, meta]) => meta.isSigner && !meta.isWritable,
+    );
+    const writableNonSigners = mapEntries.filter(
+      ([, meta]) => !meta.isSigner && meta.isWritable,
+    );
+    const readonlyNonSigners = mapEntries.filter(
+      ([, meta]) => !meta.isSigner && !meta.isWritable,
+    );
+
+    const header: MessageHeader = {
+      numRequiredSignatures: writableSigners.length + readonlySigners.length,
+      numReadonlySignedAccounts: readonlySigners.length,
+      numReadonlyUnsignedAccounts: readonlyNonSigners.length,
+    };
+
+    // sanity checks
+    {
+      assert(
+        writableSigners.length > 0,
+        'Expected at least one writable signer key',
+      );
+      const [payerAddress] = writableSigners[0];
+      assert(
+        payerAddress === this.payer.toBase58(),
+        'Expected first writable signer key to be the fee payer',
+      );
+    }
+
+    const staticAccountKeys = [
+      ...writableSigners.map(([address]) => new PublicKey(address)),
+      ...readonlySigners.map(([address]) => new PublicKey(address)),
+      ...writableNonSigners.map(([address]) => new PublicKey(address)),
+      ...readonlyNonSigners.map(([address]) => new PublicKey(address)),
+    ];
+
+    return [header, staticAccountKeys];
+  }
+
+  extractTableLookup(
+    lookupTable: AddressLookupTableAccount,
+  ): [MessageAddressTableLookup, AccountKeysFromLookups] | undefined {
+    const [writableIndexes, drainedWritableKeys] =
+      this.drainKeysFoundInLookupTable(
+        lookupTable.state.addresses,
+        keyMeta =>
+          !keyMeta.isSigner && !keyMeta.isInvoked && keyMeta.isWritable,
+      );
+    const [readonlyIndexes, drainedReadonlyKeys] =
+      this.drainKeysFoundInLookupTable(
+        lookupTable.state.addresses,
+        keyMeta =>
+          !keyMeta.isSigner && !keyMeta.isInvoked && !keyMeta.isWritable,
+      );
+
+    // Don't extract lookup if no keys were found
+    if (writableIndexes.length === 0 && readonlyIndexes.length === 0) {
+      return;
+    }
+
+    return [
+      {
+        accountKey: lookupTable.key,
+        writableIndexes,
+        readonlyIndexes,
+      },
+      {
+        writable: drainedWritableKeys,
+        readonly: drainedReadonlyKeys,
+      },
+    ];
+  }
+
+  /** @internal */
+  private drainKeysFoundInLookupTable(
+    lookupTableEntries: Array<PublicKey>,
+    keyMetaFilter: (keyMeta: CompiledKeyMeta) => boolean,
+  ): [Array<number>, Array<PublicKey>] {
+    const lookupTableIndexes = new Array();
+    const drainedKeys = new Array();
+
+    for (const [address, keyMeta] of this.keyMetaMap.entries()) {
+      if (keyMetaFilter(keyMeta)) {
+        const key = new PublicKey(address);
+        const lookupTableIndex = lookupTableEntries.findIndex(entry =>
+          entry.equals(key),
+        );
+        if (lookupTableIndex >= 0) {
+          assert(lookupTableIndex < 256, 'Max lookup table index exceeded');
+          lookupTableIndexes.push(lookupTableIndex);
+          drainedKeys.push(key);
+          this.keyMetaMap.delete(address);
+        }
+      }
+    }
+
+    return [lookupTableIndexes, drainedKeys];
+  }
+}

web3.js/src/message/index.ts

@@ -1,5 +1,7 @@
 import {PublicKey} from '../publickey';
 
+export * from './account-keys';
+// note: compiled-keys is internal and doesn't need to be exported
 export * from './legacy';
 export * from './versioned';
 export * from './v0';

web3.js/src/message/v0.ts

@@ -12,6 +12,10 @@ import {PublicKey, PUBLIC_KEY_LENGTH} from '../publickey';
 import * as shortvec from '../utils/shortvec-encoding';
 import assert from '../utils/assert';
 import {PACKET_DATA_SIZE, VERSION_PREFIX_MASK} from '../transaction/constants';
+import {TransactionInstruction} from '../transaction';
+import {AddressLookupTableAccount} from '../programs';
+import {CompiledKeys} from './compiled-keys';
+import {AccountKeysFromLookups, MessageAccountKeys} from './account-keys';
 
 /**
  * Message constructor arguments
@@ -29,6 +33,13 @@ export type MessageV0Args = {
   addressTableLookups: MessageAddressTableLookup[];
 };
 
+export type CompileV0Args = {
+  payerKey: PublicKey;
+  instructions: Array<TransactionInstruction>;
+  recentBlockhash: Blockhash;
+  addressLookupTableAccounts?: Array<AddressLookupTableAccount>;
+};
+
 export class MessageV0 {
   header: MessageHeader;
   staticAccountKeys: Array<PublicKey>;
@@ -48,6 +59,42 @@ export class MessageV0 {
     return 0;
   }
 
+  static compile(args: CompileV0Args): MessageV0 {
+    const compiledKeys = CompiledKeys.compile(args.instructions, args.payerKey);
+
+    const addressTableLookups = new Array<MessageAddressTableLookup>();
+    const accountKeysFromLookups: AccountKeysFromLookups = {
+      writable: new Array(),
+      readonly: new Array(),
+    };
+    const lookupTableAccounts = args.addressLookupTableAccounts || [];
+    for (const lookupTable of lookupTableAccounts) {
+      const extractResult = compiledKeys.extractTableLookup(lookupTable);
+      if (extractResult !== undefined) {
+        const [addressTableLookup, {writable, readonly}] = extractResult;
+        addressTableLookups.push(addressTableLookup);
+        accountKeysFromLookups.writable.push(...writable);
+        accountKeysFromLookups.readonly.push(...readonly);
+      }
+    }
+
+    const [header, staticAccountKeys] = compiledKeys.getMessageComponents();
+    const accountKeys = new MessageAccountKeys(
+      staticAccountKeys,
+      accountKeysFromLookups,
+    );
+    const compiledInstructions = accountKeys.compileInstructions(
+      args.instructions,
+    );
+    return new MessageV0({
+      header,
+      staticAccountKeys,
+      recentBlockhash: args.recentBlockhash,
+      compiledInstructions,
+      addressTableLookups,
+    });
+  }
+
   serialize(): Uint8Array {
     const encodedStaticAccountKeysLength = Array<number>();
     shortvec.encodeLength(

web3.js/src/publickey.ts

@@ -40,6 +40,9 @@ function isPublicKeyData(value: PublicKeyInitData): value is PublicKeyData {
   return (value as PublicKeyData)._bn !== undefined;
 }
 
+// local counter used by PublicKey.unique()
+let uniquePublicKeyCounter = 1;
+
 /**
  * A public key
  */
@@ -73,6 +76,15 @@ export class PublicKey extends Struct {
     }
   }
 
+  /**
+   * Returns a unique PublicKey for tests and benchmarks using acounter
+   */
+  static unique(): PublicKey {
+    const key = new PublicKey(uniquePublicKeyCounter);
+    uniquePublicKeyCounter += 1;
+    return new PublicKey(key.toBuffer());
+  }
+
   /**
    * Default public key value. (All zeros)
    */