feat: handle mutex calls while prepare resolved ops

packages/cwait/src/onchain/context/context.ts

@@ -1,16 +1,39 @@
-import { Context, extractUser, getMethodArguments } from '@coinweb/contract-kit';
+import {
+  Context,
+  ContractIssuer,
+  extractContractArgs,
+  extractUser,
+  getMethodArguments,
+  GTake,
+  User,
+} from '@coinweb/contract-kit';
+import { CwebTake } from '@coinweb/contract-kit/dist/types/operations/take';
 import { getCallParameters, getContractIssuer } from 'lib/onchain';
 
-import { ExecutorMethodArgs } from '../../types';
+import { pushAwaitedOp } from '../../../dist/onchain/ops';
+import { ExecutorMethodArgs, ResolvedOp } from '../../types';
 import { uuid } from '../utils';
 
-import { extractFunds } from './extractFunds';
 import { extractOps } from './extractOps';
 
 let rawContext: Context | null = null;
 
-export const setRawContext = (ctx: Context) => {
-  rawContext = ctx;
+const initialContext = {
+  isChild: false,
+  thisId: '',
+  parentId: '' as string | undefined,
+  ops: [] as ResolvedOp[],
+  funds: {
+    availableCweb: 0n,
+    storedCweb: 0n,
+    takeOps: [] as GTake<CwebTake>[],
+  },
+  methodName: '',
+  initialArgs: [] as unknown[],
+  user: null as User | null,
+  issuer: null as ContractIssuer | null,
+  parentTxId: '',
+  needSaveResult: false,
 };
 
 export const getRawContext = () => {
@@ -25,71 +48,106 @@ export const getMethodArgs = () => {
   return getMethodArguments(getRawContext()) as [methodName: string, ...ExecutorMethodArgs];
 };
 
-let thisId: string | undefined;
+export const handleContext = (ctx: Context) => {
+  rawContext = ctx;
+
+  const contractArgs = extractContractArgs(getRawContext().tx);
+  const [
+    methodName,
+    initialArgs,
+    resolvedOps,
+    caller,
+    thisId,
+    parentId,
+    saveResult,
+    takenFundsIds = [],
+    execOpsIndexes = [],
+  ] = getMethodArgs();
+
+  initialContext.isChild = !!parentId;
+  initialContext.thisId = thisId ?? uuid();
+  initialContext.parentId = parentId;
+  initialContext.methodName = methodName;
+  initialContext.initialArgs = initialArgs ?? [];
+  initialContext.needSaveResult = saveResult ?? false;
+
+  const { authInfo } = getCallParameters(getRawContext());
+  initialContext.user = (authInfo && extractUser(authInfo)) ?? caller ?? null;
+
+  initialContext.issuer = getContractIssuer(getRawContext());
+  initialContext.parentTxId = getRawContext().call.txid;
 
-export const context = {
-  get isChild() {
-    return !!getMethodArgs()[5];
-  },
-  get thisId() {
-    thisId = thisId ?? getMethodArgs()[4] ?? uuid();
-    return thisId;
-  },
-  get parentId() {
-    return getMethodArgs()[5];
-  },
-  get ops() {
   console.log('Parse ops in context');
+  const { extractedOps, executionOpsTakeOp, storedCweb, fundsTakeOps } = extractOps({
+    contractArgs,
+    takenFundsIds,
+    execOpsIndexes,
+  });
 
-    const previousOps = getMethodArgs()[2] ?? [];
-    const { extractedOps } = extractOps();
-
+  const previousOps = resolvedOps ?? [];
   const allResolvedOps = [...previousOps, ...extractedOps];
-
   console.log('new ops >>>', JSON.stringify(extractedOps));
+  initialContext.ops = allResolvedOps;
 
-    return allResolvedOps;
-  },
-  get funds() {
   const { availableCweb } = getCallParameters(getRawContext());
 
-    const { amount: storedCweb, takeOps } = extractFunds();
-
   console.log('Available Cweb: ', availableCweb);
   console.log('Stored Cweb: ', storedCweb);
-    console.log('Take Ops: ', takeOps);
+  console.log('Take Ops: ', fundsTakeOps);
 
-    return {
-      availableCweb,
-      storedCweb,
-      takeOps,
+  initialContext.funds = {
+    availableCweb: availableCweb,
+    storedCweb: storedCweb,
+    takeOps: fundsTakeOps,
   };
+
+  if (executionOpsTakeOp) {
+    pushAwaitedOp(executionOpsTakeOp);
+  }
+
+  return !!executionOpsTakeOp; //This is a flag to check if the execution should restart / TODO: refactor;
+};
+
+export const context = {
+  get isChild() {
+    return initialContext.isChild;
+  },
+  get thisId() {
+    return initialContext.thisId;
+  },
+  get parentId() {
+    return initialContext.parentId;
   },
   get methodName() {
     return getMethodArguments(getRawContext())[0] as string;
   },
   get initialArgs() {
-    return getMethodArgs()[1] ?? [];
+    return initialContext.initialArgs;
   },
   get user() {
-    const { authInfo } = getCallParameters(getRawContext());
-    const provided = getMethodArgs()[3];
-
-    const user = (authInfo && extractUser(authInfo)) ?? provided;
-
-    if (!user) {
+    if (!initialContext.user) {
       throw new Error('User not found');
     }
 
-    return user;
+    return initialContext.user;
   },
   get issuer() {
-    return getContractIssuer(getRawContext());
+    if (!initialContext.issuer) {
+      throw new Error('Issuer not found');
+    }
+
+    return initialContext.issuer;
   },
   get parentTxId() {
-    return getRawContext().call.txid;
+    return initialContext.parentTxId;
+  },
+  get ops() {
+    return initialContext.ops;
   },
   get needSaveResult() {
-    return getMethodArgs()[6] ?? false;
+    return initialContext.needSaveResult;
+  },
+  get funds() {
+    return initialContext.funds;
   },
 };

packages/cwait/src/onchain/context/extractFunds.ts

@@ -1,34 +0,0 @@
-import { constructTake } from '@coinweb/contract-kit';
-import { extractRead } from '@coinweb/contract-kit/dist/esm/operations/read';
-
-import { getMethodArgs } from './context';
-import { extractOps } from './extractOps';
-
-export const extractFunds = () => {
-  const { fundsOp } = extractOps();
-
-  const takenFundsIds = getMethodArgs()[7] ?? [];
-
-  if (!fundsOp) {
-    return {
-      takeOps: [],
-      amount: 0n,
-    };
-  }
-
-  const claims = extractRead(fundsOp);
-
-  const newClaims = claims?.filter(({ content }) => !takenFundsIds.includes((content.key.second_part as [string])[0]));
-
-  if (!newClaims?.length) {
-    return {
-      takeOps: [],
-      amount: 0n,
-    };
-  }
-
-  return {
-    takeOps: newClaims.map(({ content }) => constructTake(content.key)),
-    amount: newClaims.reduce((acc, { content }) => acc + BigInt(content.fees_stored), 0n),
-  };
-};

packages/cwait/src/onchain/context/extractOps.ts

@@ -1,18 +1,51 @@
-import { extractContractArgs } from '@coinweb/contract-kit';
+import { constructTake, extractRead, GRead, GTake, ResolvedOperation, ResolvedRead } from '@coinweb/contract-kit';
+import { CwebTake } from '@coinweb/contract-kit/dist/types/operations/take';
 
 import { ResolvedOp, TypedClaim } from '../../types';
 import { resultKey } from '../claims/result';
+import { mutexBlockLockKey, mutexBlockUnlockKey, mutexExecOpsKey, MutexExecOpsResult } from '../mutex';
 import { isResolvedBlockOp, isResolvedReadOp, isResolvedTakeOp } from '../utils';
 
-import { getRawContext } from './context';
+const extractFunds = (fundsOp: GRead<ResolvedRead>, takenFundsIds: string[]) => {
+  if (!fundsOp) {
+    return {
+      takeOps: [],
+      amount: 0n,
+    };
+  }
 
-export const extractOps = () => {
-  const contractArgs = extractContractArgs(getRawContext().tx);
+  const claims = extractRead(fundsOp);
 
+  const newClaims = claims?.filter(({ content }) => !takenFundsIds.includes((content.key.second_part as [string])[0]));
+
+  if (!newClaims?.length) {
+    return {
+      takeOps: [],
+      amount: 0n,
+    };
+  }
+
+  return {
+    takeOps: newClaims.map(({ content }) => constructTake(content.key)),
+    amount: newClaims.reduce((acc, { content }) => acc + BigInt(content.fees_stored), 0n),
+  };
+};
+
+export const extractOps = ({
+  contractArgs,
+  takenFundsIds,
+  execOpsIndexes,
+}: {
+  contractArgs: ResolvedOperation[];
+  takenFundsIds: string[];
+  execOpsIndexes: number[];
+}) => {
   if (!contractArgs.length) {
     return {
       extractedOps: [],
-      fundsOps: [],
+      fundsTakeOps: [],
+      storedCweb: 0n,
+      executionOpsTakeOp: null,
     };
   }
 
@@ -23,6 +56,9 @@ export const extractOps = () => {
   }
 
   const extractedOps: ResolvedOp[] = [];
+  let executedOps: MutexExecOpsResult = [];
+  let executionOpsProfit: bigint = 0n;
+  let executionOpsTakeOp: GTake<CwebTake> | null = null;
 
   let i = 0;
 
@@ -34,6 +70,8 @@ export const extractOps = () => {
 
       //Maybe it is needed to check more conditions here
       if (Array.isArray(first) && first[0] === resultKey) {
+        switch (first[0]) {
+          case resultKey: {
             const nextAfterBlock = resolvedOps[i + 1];
 
             if (!isResolvedTakeOp(nextAfterBlock)) {
@@ -45,14 +83,96 @@ export const extractOps = () => {
             i += 2;
             continue;
           }
+          case mutexBlockLockKey: {
+            const nextAfterBlock = resolvedOps[i + 1];
+
+            if (!isResolvedTakeOp(nextAfterBlock)) {
+              throw new Error('Wrong mutex lock result');
+            }
+
+            extractedOps.push({ SlotOp: { ok: true } });
+
+            i += 2;
+            continue;
+          }
+          case mutexBlockUnlockKey: {
+            const nextAfterBlock = resolvedOps[i + 1];
+
+            if (!isResolvedTakeOp(nextAfterBlock)) {
+              throw new Error('Wrong mutex unlock result');
+            }
+
+            extractedOps.push({ SlotOp: { ok: true } });
+
+            i += 2;
+            continue;
+          }
+          case mutexExecOpsKey: {
+            const nextAfterBlock = resolvedOps[i + 1];
+
+            const execOpResultClaim = extractRead(nextAfterBlock)?.[0]?.content;
+
+            if (!execOpResultClaim) {
+              throw new Error('Wrong mutex exec result');
+            }
+
+            executedOps = execOpResultClaim.body as MutexExecOpsResult;
+            executionOpsProfit = BigInt(execOpResultClaim.fees_stored);
+            executionOpsTakeOp = constructTake(execOpResultClaim.key);
+            i += 2;
+            continue;
+          }
+
+          default:
+            break;
+        }
+      }
+    }
+
+    if (isResolvedTakeOp(op)) {
+      const keyFirstPart = op.TakeOp.result.key.first_part;
+      if (Array.isArray(keyFirstPart) && keyFirstPart[0] === mutexExecOpsKey) {
+        i++;
+        continue;
+      }
     }
 
     extractedOps.push(op);
     i++;
   }
 
+  const allOps: ResolvedOp[] = [];
+
+  for (let i = 0; i < resolvedOps.length + executedOps.length; i++) {
+    if (execOpsIndexes.includes(i)) {
+      const op = executedOps.shift();
+
+      if (!op) {
+        throw new Error('Wrong mutex exec result');
+      }
+
+      if (op.ok) {
+        allOps.push(op.resolved);
+      } else {
+        allOps.push({ SlotOp: { ok: false, error: op.error } });
+      }
+    } else {
+      const op = resolvedOps.shift();
+
+      if (!op) {
+        throw new Error('Contract call args parsing error');
+      }
+
+      allOps.push(op);
+    }
+  }
+
+  const { takeOps, amount } = extractFunds(fundsOp, takenFundsIds);
+
   return {
-    fundsOp,
-    extractedOps,
+    fundsTakeOps: takeOps,
+    storedCweb: amount + executionOpsProfit,
+    extractedOps: allOps,
+    executionOpsTakeOp,
   };
 };

packages/cwait/src/onchain/executor.ts

@@ -1,7 +1,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { Context, NewTx, getMethodArguments, isSelfCall } from '@coinweb/contract-kit';
 
-import { context, getRawContext, setRawContext } from './context';
+import { context, getRawContext, handleContext } from './context';
 import { pushResolvedOp } from './promisifiedOps/resolved';
 import { constructTx } from './utils';
 
@@ -10,13 +10,17 @@ let abortExecution: ((result: boolean) => void) | null = null;
 export const executor = (method: (...args: any[]) => Promise<void>) => {
   return async (ctx: Context): Promise<NewTx[]> => {
     console.log('executor-start >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>');
-    setRawContext(ctx);
+    const shouldRestart = handleContext(ctx);
     pushResolvedOp(context.ops);
 
     if (getMethodArguments(getRawContext()).length > 2 && !isSelfCall(ctx)) {
       throw new Error('Wrong contract call, check the call arguments');
     }
 
+    if (shouldRestart) {
+      return constructTx(false);
+    }
+
     const execution = new Promise<boolean>((resolve, reject) => {
       abortExecution = resolve;
 

packages/cwait/src/onchain/mutex/calls/constructExecOpsCall.ts

@@ -4,7 +4,7 @@ import { CwebTake } from '@coinweb/contract-kit/dist/types/operations/take';
 import { PreparedExtendedStoreOp } from '../../../types';
 import {
   constructMutexExecOpsBlock,
-  constructMutexExecOpsClaimTake,
+  constructMutexExecOpsClaimRead,
   constructMutexLockClaimRangeRead,
 } from '../claims';
 import { execOpsMethodName } from '../methods';
@@ -65,6 +65,8 @@ export const constructExecOpsCall = ({
       contractArgs: [constructMutexLockClaimRangeRead(issuer)],
     },
     fee: providedCweb + contractArgsFee,
-    ops: execId ? ([constructMutexExecOpsBlock(execId, issuer), constructMutexExecOpsClaimTake(execId)] as const) : [],
+    ops: execId
+      ? ([constructMutexExecOpsBlock(execId, issuer), constructMutexExecOpsClaimRead(issuer, execId)] as const)
+      : [],
   };
 };

packages/cwait/src/onchain/mutex/claims/mutexExecOps.ts

@@ -3,6 +3,7 @@ import {
   constructBlock,
   constructClaim,
   constructClaimKey,
+  constructRead,
   constructStore,
   constructTake,
   ContractIssuer,
@@ -26,6 +27,9 @@ export const constructMutexExecOpsClaimStore = (
   storeCweb: bigint = 0n
 ): GStore<CwebStore> => constructStore(constructMutexExecOpsClaim(execId, result, storeCweb));
 
+export const constructMutexExecOpsClaimRead = (issuer: ContractIssuer, execId: string) =>
+  constructRead(issuer, constructMutexExecOpsClaimKey(execId));
+
 export const constructMutexExecOpsClaimTake = (execId: string) => constructTake(constructMutexExecOpsClaimKey(execId));
 
 export const constructMutexExecOpsFilter = (lockId: string, issuer: ContractIssuer): BlockFilter => {

packages/cwait/src/onchain/mutex/index.ts

@@ -10,8 +10,10 @@ import {
   saveExecOpResult,
   saveExecOpResultMethodName,
   unlockMethodName,
+  waitMethodName,
 } from './methods';
 import { preReadExecTakeOps } from './methods/preReadExecTakeOps';
+import { wait } from './methods/wait';
 
 export * from './claims';
 export * from './types';
@@ -24,4 +26,5 @@ export const mutexMethods = {
   [execOpsMethodName]: execOps,
   [preReadExecTakeOpsMethodName]: preReadExecTakeOps,
   [saveExecOpResultMethodName]: saveExecOpResult,
+  [waitMethodName]: wait,
 };

packages/cwait/src/onchain/mutex/methods/execLock.ts

@@ -13,7 +13,7 @@ import { getCallParameters, getContractIssuer } from 'lib/onchain';
 import { toHex, TypedClaim } from 'lib/shared';
 
 import { lockFee } from '../settings';
-import { MutexLockState } from '../types';
+import { MutexLockState, MutexWaitNotifyArgs } from '../types';
 import { isMatchLockKeys } from '../utils';
 
 import { lockMethodName, notifyLockMethodName } from './names';
@@ -74,7 +74,7 @@ export const mutexExecLock = (context: Context) => {
             ref: constructContractRef(issuer, []),
             methodInfo: {
               methodName: notifyLockMethodName,
-              methodArgs: [(lockCandidate.key.second_part as [number, string])[1]],
+              methodArgs: [0, (lockCandidate.key.second_part as [number, string])[1]] satisfies MutexWaitNotifyArgs,
             },
             contractInfo: {
               providedCweb: 200n,

packages/cwait/src/onchain/mutex/methods/names.ts

@@ -5,3 +5,4 @@ export const notifyLockMethodName = '_mutex_notify_lock';
 export const execOpsMethodName = '_mutex_execOps';
 export const saveExecOpResultMethodName = '_mutex_save_exec_op_result';
 export const preReadExecTakeOpsMethodName = '_mutex_pre_read_exec_take_ops';
+export const waitMethodName = '_mutex_wait';

packages/cwait/src/onchain/mutex/methods/wait.ts

@@ -0,0 +1,61 @@
+import { constructContinueTx, constructContractRef, Context } from '@coinweb/contract-kit';
+import { getContractArguments, getContractIssuer } from 'lib/onchain';
+
+import { waitSteps } from '../settings';
+import { MutexNotifyLockArgs, MutexWaitNotifyArgs } from '../types';
+
+import { notifyLockMethodName } from './names';
+
+export const wait = (context: Context) => {
+  const issuer = getContractIssuer(context);
+  const [step, lockId] = getContractArguments<MutexWaitNotifyArgs>(context);
+
+  //Including the notifyLock method as step
+  if (step + 1 < waitSteps) {
+    return [
+      constructContinueTx(
+        context,
+        [],
+        [
+          {
+            callInfo: {
+              ref: constructContractRef(issuer, []),
+              methodInfo: {
+                methodName: notifyLockMethodName,
+                methodArgs: [step + 1, lockId] satisfies MutexWaitNotifyArgs,
+              },
+              contractInfo: {
+                providedCweb: 200n,
+                authenticated: null,
+              },
+              contractArgs: [],
+            },
+          },
+        ]
+      ),
+    ];
+  }
+
+  return [
+    constructContinueTx(
+      context,
+      [],
+      [
+        {
+          callInfo: {
+            ref: constructContractRef(issuer, []),
+            methodInfo: {
+              methodName: notifyLockMethodName,
+              methodArgs: [lockId] satisfies MutexNotifyLockArgs,
+            },
+            contractInfo: {
+              providedCweb: 200n,
+              authenticated: null,
+            },
+            contractArgs: [],
+          },
+        },
+      ]
+    ),
+  ];
+};

packages/cwait/src/onchain/mutex/settings.ts

@@ -2,3 +2,5 @@ export const lockFee = 5000n;
 export const unlockFee = 1000n + lockFee;
 
 export const rangeReadLimit = 10000;
+
+export const waitSteps = 2;

packages/cwait/src/onchain/mutex/types.ts

@@ -24,6 +24,8 @@ export type MutexUnlockArgs = [lockId: string, timestamp: number, notify?: boole
 
 export type MutexExecOpArgs = [ops: (GTake<CwebTake> | GStore<CwebStore>)[], processId: string, execId?: string];
 
+export type MutexWaitNotifyArgs = [step: number, lockId: string];
+
 export type MutexNotifyLockArgs = [lockId: string];
 
 export type MutexPreReadTakeOpsArgs = [execId: string, unavailableIndexes: number[], ops: PreparedOperation[]];