Made baselayer contracts more extensible using inheritance

Author: rpanic

packages/protocol/src/settlement/contracts/DispatchContractProtocolModule.ts

@@ -9,6 +9,7 @@ import {
 import {
   DispatchSmartContract,
   DispatchContractType,
+  DispatchSmartContractBase,
 } from "./DispatchSmartContract";
 
 export type DispatchContractConfig = {
@@ -29,7 +30,7 @@ export class DispatchContractProtocolModule extends ContractModule<
     const { incomingMessagesMethods } = this.config;
     const methodIdMappings = this.runtime.methodIdResolver.methodIdMap();
 
-    DispatchSmartContract.args = {
+    DispatchSmartContractBase.args = {
       incomingMessagesPaths: incomingMessagesMethods,
       methodIdMappings,
     };

packages/protocol/src/settlement/contracts/DispatchSmartContract.ts

@@ -32,23 +32,19 @@ export interface DispatchContractType {
   promisedMessagesHash: State<Field>;
 }
 
-export class DispatchSmartContract
-  extends SmartContract
-  implements DispatchContractType
-{
+export abstract class DispatchSmartContractBase extends SmartContract {
   public static args: {
     methodIdMappings: RuntimeMethodIdMapping;
     incomingMessagesPaths: Record<string, `${string}.${string}`>;
   };
 
-  @state(Field) public promisedMessagesHash = State<Field>();
+  abstract promisedMessagesHash: State<Field>;
 
-  @state(Field) public honoredMessagesHash = State<Field>();
+  abstract honoredMessagesHash: State<Field>;
 
-  @state(PublicKey) public settlementContract = State<PublicKey>();
+  abstract settlementContract: State<PublicKey>;
 
-  @method
-  public async updateMessagesHash(
+  protected async updateMessagesHashBase(
     executedMessagesHash: Field,
     newPromisedMessagesHash: Field
   ) {
@@ -65,8 +61,7 @@ export class DispatchSmartContract
     this.promisedMessagesHash.set(newPromisedMessagesHash);
   }
 
-  @method
-  public async initialize(settlementContract: PublicKey) {
+  protected async initializeBase(settlementContract: PublicKey) {
     this.promisedMessagesHash.getAndRequireEquals().assertEquals(Field(0));
     this.honoredMessagesHash.getAndRequireEquals().assertEquals(Field(0));
     this.settlementContract
@@ -78,7 +73,7 @@ export class DispatchSmartContract
     this.settlementContract.set(settlementContract);
   }
 
-  private dispatchMessage<Type>(
+  protected dispatchMessage<Type>(
     methodId: Field,
     value: Type,
     valueType: ProvableExtended<Type>
@@ -106,6 +101,33 @@ export class DispatchSmartContract
       data: [args],
     };
   }
+}
+
+export class DispatchSmartContract
+  extends DispatchSmartContractBase
+  implements DispatchContractType
+{
+  @state(Field) public promisedMessagesHash = State<Field>();
+
+  @state(Field) public honoredMessagesHash = State<Field>();
+
+  @state(PublicKey) public settlementContract = State<PublicKey>();
+
+  @method
+  public async updateMessagesHash(
+    executedMessagesHash: Field,
+    newPromisedMessagesHash: Field
+  ) {
+    return await this.updateMessagesHashBase(
+      executedMessagesHash,
+      newPromisedMessagesHash
+    );
+  }
+
+  @method
+  public async initialize(settlementContract: PublicKey) {
+    return await this.initializeBase(settlementContract);
+  }
 
   @method
   public async deposit(amount: UInt64) {
@@ -126,7 +148,7 @@ export class DispatchSmartContract
     });
 
     const { methodIdMappings, incomingMessagesPaths } =
-      DispatchSmartContract.args;
+      DispatchSmartContractBase.args;
 
     const methodId = Field(
       methodIdMappings[incomingMessagesPaths.deposit].methodId

packages/protocol/src/settlement/contracts/SettlementContractProtocolModule.ts

@@ -14,6 +14,7 @@ import {
   LazyBlockProof,
   SettlementContractType,
   SettlementSmartContract,
+  SettlementSmartContractBase,
 } from "./SettlementSmartContract";
 
 export type SettlementContractConfig = {
@@ -55,7 +56,7 @@ export class SettlementContractProtocolModule extends ContractModule<
     const escapeHatchSlotsInterval =
       config.escapeHatchSlotsInterval ?? DEFAULT_ESCAPE_HATCH;
 
-    SettlementSmartContract.args = {
+    SettlementSmartContractBase.args = {
       DispatchContract: dispatchContract,
       hooks,
       withdrawalStatePath: withdrawalStatePathSplit,

packages/protocol/src/settlement/contracts/SettlementSmartContract.ts

@@ -83,10 +83,7 @@ export interface SettlementContractType {
 // Some random prefix for the sequencer signature
 export const BATCH_SIGNATURE_PREFIX = prefixToField("pk-batchSignature");
 
-export class SettlementSmartContract
-  extends TokenContract
-  implements SettlementContractType
-{
+export abstract class SettlementSmartContractBase extends TokenContract {
   // This pattern of injecting args into a smartcontract is currently the only
   // viable solution that works given the inheritance issues of o1js
   public static args: {
@@ -96,23 +93,18 @@ export class SettlementSmartContract
     escapeHatchSlotsInterval: number;
   };
 
-  @state(Field) public sequencerKey = State<Field>();
-  @state(UInt32) public lastSettlementL1BlockHeight = State<UInt32>();
-
-  @state(Field) public stateRoot = State<Field>();
-  @state(Field) public networkStateHash = State<Field>();
-  @state(Field) public blockHashRoot = State<Field>();
-
-  @state(Field) public dispatchContractAddressX = State<Field>();
-
-  @state(Field) public outgoingMessageCursor = State<Field>();
+  abstract sequencerKey: State<Field>;
+  abstract lastSettlementL1BlockHeight: State<UInt32>;
+  abstract stateRoot: State<Field>;
+  abstract networkStateHash: State<Field>;
+  abstract blockHashRoot: State<Field>;
+  abstract dispatchContractAddressX: State<Field>;
+  abstract outgoingMessageCursor: State<Field>;
 
-  @method async approveBase(forest: AccountUpdateForest) {
-    this.checkZeroBalanceChange(forest);
-  }
-
-  @method
-  public async initialize(sequencer: PublicKey, dispatchContract: PublicKey) {
+  protected async initializeBase(
+    sequencer: PublicKey,
+    dispatchContract: PublicKey
+  ) {
     this.sequencerKey.getAndRequireEquals().assertEquals(Field(0));
     this.stateRoot.getAndRequireEquals().assertEquals(Field(0));
     this.blockHashRoot.getAndRequireEquals().assertEquals(Field(0));
@@ -125,13 +117,12 @@ export class SettlementSmartContract
     this.networkStateHash.set(NetworkState.empty().hash());
     this.dispatchContractAddressX.set(dispatchContract.x);
 
-    const { DispatchContract } = SettlementSmartContract.args;
+    const { DispatchContract } = SettlementSmartContractBase.args;
     const contractInstance = new DispatchContract(dispatchContract);
     await contractInstance.initialize(this.address);
   }
 
-  @method
-  public async settle(
+  protected async settleBase(
     blockProof: LazyBlockProof,
     signature: Signature,
     dispatchContractAddress: PublicKey,
@@ -159,7 +150,7 @@ export class SettlementSmartContract
     );
 
     const { DispatchContract, escapeHatchSlotsInterval, hooks } =
-      SettlementSmartContract.args;
+      SettlementSmartContractBase.args;
 
     // Get dispatch contract values
     // These values are witnesses but will be checked later on the AU
@@ -265,14 +256,60 @@ export class SettlementSmartContract
 
     this.lastSettlementL1BlockHeight.set(minBlockHeightIncluded);
   }
+}
+
+export class SettlementSmartContract
+  extends SettlementSmartContractBase
+  implements SettlementContractType
+{
+  @state(Field) public sequencerKey = State<Field>();
+  @state(UInt32) public lastSettlementL1BlockHeight = State<UInt32>();
+
+  @state(Field) public stateRoot = State<Field>();
+  @state(Field) public networkStateHash = State<Field>();
+  @state(Field) public blockHashRoot = State<Field>();
+
+  @state(Field) public dispatchContractAddressX = State<Field>();
+
+  @state(Field) public outgoingMessageCursor = State<Field>();
+
+  @method async approveBase(forest: AccountUpdateForest) {
+    this.checkZeroBalanceChange(forest);
+  }
+
+  @method
+  public async initialize(sequencer: PublicKey, dispatchContract: PublicKey) {
+    return await this.initializeBase(sequencer, dispatchContract);
+  }
+
+  @method
+  public async settle(
+    blockProof: LazyBlockProof,
+    signature: Signature,
+    dispatchContractAddress: PublicKey,
+    publicKey: PublicKey,
+    inputNetworkState: NetworkState,
+    outputNetworkState: NetworkState,
+    newPromisedMessagesHash: Field
+  ) {
+    return await this.settleBase(
+      blockProof,
+      signature,
+      dispatchContractAddress,
+      publicKey,
+      inputNetworkState,
+      outputNetworkState,
+      newPromisedMessagesHash
+    );
+  }
 
   @method
   public async rollupOutgoingMessages(batch: OutgoingMessageArgumentBatch) {
     let counter = this.outgoingMessageCursor.getAndRequireEquals();
     const stateRoot = this.stateRoot.getAndRequireEquals();
 
     const [withdrawalModule, withdrawalStateName] =
-      SettlementSmartContract.args.withdrawalStatePath;
+      SettlementSmartContractBase.args.withdrawalStatePath;
     const mapPath = Path.fromProperty(withdrawalModule, withdrawalStateName);
 
     let accountCreationFeePaid = Field(0);

packages/protocol/src/settlement/modularity/ProvableSettlementHook.ts

@@ -3,7 +3,7 @@ import { Field, PublicKey, UInt32 } from "o1js";
 import { ProtocolModule } from "../../protocol/ProtocolModule";
 import { NetworkState } from "../../model/network/NetworkState";
 import type { BlockProof } from "../../prover/block/BlockProver";
-import type { SettlementSmartContract } from "../contracts/SettlementSmartContract";
+import type { SettlementSmartContractBase } from "../contracts/SettlementSmartContract";
 
 export type SettlementStateRecord = {
   sequencerKey: PublicKey;
@@ -27,7 +27,7 @@ export abstract class ProvableSettlementHook<
   Config,
 > extends ProtocolModule<Config> {
   public abstract beforeSettlement(
-    smartContract: SettlementSmartContract,
+    smartContract: SettlementSmartContractBase,
     inputs: SettlementHookInputs
   ): Promise<void>;
 }