Omit unused args when calling "use cache" functions (#72506)

Using the information byte in a server reference ID, as introduced in
#71463, we can omit unused arguments when calling a `"use cache"`
function. This enables us, for example, to use a cached function in
conjunction with `useActionState` when the function does not need the
previous state, as shown below:

```ts
'use cache'

export async function getCachedRandomValue() {
  return Math.random()
}
```

```tsx
'use client'

import { useActionState } from 'react'
import { getCachedRandomValue } from './cached'

export default function Page() {
  const [result, formAction] = useActionState(getCachedRandomValue, null)

  return (
    <form action={formAction}>
      <p>{result}</p>
      <button>Submit</button>
    </form>
  )
}
```

Previously, in this use case, a new value would have been rendered for
every click. This is because React passes the previous state and the
form data as arguments into the action function. These two arguments are
then included in the cache key, resulting in a new cache key being
generated for every invocation.

_Note: This optimization is currently based purely on the function
signature. It does not yet detect when a param is declared, but not used
in the function body. It's also not yet applied for server actions._

packages/next/src/client/components/router-reducer/reducers/server-action-reducer.ts

@@ -51,6 +51,10 @@ import { getRedirectError, RedirectType } from '../../redirect'
 import { createSeededPrefetchCacheEntry } from '../prefetch-cache-utils'
 import { removeBasePath } from '../../../remove-base-path'
 import { hasBasePath } from '../../../has-base-path'
+import {
+  extractInfoFromServerReferenceId,
+  omitUnusedArgs,
+} from './server-reference-info'
 
 type FetchServerActionResult = {
   redirectLocation: URL | undefined
@@ -71,7 +75,15 @@ async function fetchServerAction(
   { actionId, actionArgs }: ServerActionAction
 ): Promise<FetchServerActionResult> {
   const temporaryReferences = createTemporaryReferenceSet()
-  const body = await encodeReply(actionArgs, { temporaryReferences })
+  const info = extractInfoFromServerReferenceId(actionId)
+
+  // TODO: Currently, we're only omitting unused args for the experimental "use
+  // cache" functions. Once the server reference info byte feature is stable, we
+  // should apply this to server actions as well.
+  const usedArgs =
+    info.type === 'use-cache' ? omitUnusedArgs(actionArgs, info) : actionArgs
+
+  const body = await encodeReply(usedArgs, { temporaryReferences })
 
   const res = await fetch('', {
     method: 'POST',

packages/next/src/client/components/router-reducer/reducers/server-reference-info.test.ts

@@ -0,0 +1,193 @@
+import {
+  type ServerReferenceInfo,
+  extractInfoFromServerReferenceId,
+  omitUnusedArgs,
+} from './server-reference-info'
+
+describe('extractInfoFromServerReferenceId', () => {
+  test('should parse id with typeBit 0, no args used, no restArgs', () => {
+    const id = '00' // 0b00000000
+
+    const expected: ServerReferenceInfo = {
+      type: 'server-action',
+      usedArgs: [false, false, false, false, false, false],
+      hasRestArgs: false,
+    }
+
+    expect(extractInfoFromServerReferenceId(id)).toEqual(expected)
+  })
+
+  test('should parse id with typeBit 1, all args used, restArgs true', () => {
+    const id = 'ff' // 0b11111111
+
+    const expected: ServerReferenceInfo = {
+      type: 'use-cache',
+      usedArgs: [true, true, true, true, true, true],
+      hasRestArgs: true,
+    }
+
+    expect(extractInfoFromServerReferenceId(id)).toEqual(expected)
+  })
+
+  test('should parse id with typeBit 0, argMask 0b101010, restArgs false', () => {
+    const id = '54' // 0b01010100
+
+    const expected: ServerReferenceInfo = {
+      type: 'server-action',
+      usedArgs: [true, false, true, false, true, false],
+      hasRestArgs: false,
+    }
+
+    expect(extractInfoFromServerReferenceId(id)).toEqual(expected)
+  })
+
+  test('should parse id with typeBit 1, argMask 0b000101, restArgs true', () => {
+    const id = '8b' // 0b10001011
+
+    const expected: ServerReferenceInfo = {
+      type: 'use-cache',
+      usedArgs: [false, false, false, true, false, true],
+      hasRestArgs: true,
+    }
+
+    expect(extractInfoFromServerReferenceId(id)).toEqual(expected)
+  })
+})
+
+describe('omitUnusedArgs', () => {
+  test('should return empty array when no args are used and no restArgs', () => {
+    const args = ['arg1', 'arg2', 'arg3']
+
+    const info: ServerReferenceInfo = {
+      type: 'server-action',
+      usedArgs: [false, false, false, false, false, false],
+      hasRestArgs: false,
+    }
+
+    expect(omitUnusedArgs(args, info)).toEqual([])
+  })
+
+  test('should return all args when all args are used and has restArgs', () => {
+    const args = [
+      'arg1',
+      'arg2',
+      'arg3',
+      'arg4',
+      'arg5',
+      'arg6',
+      'restArg1',
+      'restArg2',
+    ]
+
+    const info: ServerReferenceInfo = {
+      type: 'use-cache',
+      usedArgs: [true, true, true, true, true, true],
+      hasRestArgs: true,
+    }
+
+    expect(omitUnusedArgs(args, info)).toEqual(args)
+  })
+
+  test('should filter args when some args are used and no restArgs', () => {
+    const args = ['arg1', 'arg2', 'arg3', 'arg4', 'arg5', 'arg6']
+
+    const info: ServerReferenceInfo = {
+      type: 'server-action',
+      usedArgs: [true, false, true, false, true, false],
+      hasRestArgs: false,
+    }
+
+    expect(omitUnusedArgs(args, info)).toEqual([
+      'arg1',
+      undefined,
+      'arg3',
+      undefined,
+      'arg5',
+      undefined,
+    ])
+  })
+
+  test('should include restArgs when hasRestArgs is true', () => {
+    const args = [
+      'arg1',
+      'arg2',
+      'arg3',
+      'arg4',
+      'arg5',
+      'arg6',
+      'restArg1',
+      'restArg2',
+    ]
+
+    const info: ServerReferenceInfo = {
+      type: 'use-cache',
+      usedArgs: [false, false, false, true, false, true],
+      hasRestArgs: true,
+    }
+
+    expect(omitUnusedArgs(args, info)).toEqual([
+      undefined,
+      undefined,
+      undefined,
+      'arg4',
+      undefined,
+      'arg6',
+      'restArg1',
+      'restArg2',
+    ])
+  })
+
+  test('should not include extra args when hasRestArgs is false', () => {
+    const args = [
+      'arg1',
+      'arg2',
+      'arg3',
+      'arg4',
+      'arg5',
+      'arg6',
+      'extraArg1',
+      'extraArg2',
+    ]
+
+    const info: ServerReferenceInfo = {
+      type: 'server-action',
+      usedArgs: [true, true, true, true, true, true],
+      hasRestArgs: false,
+    }
+
+    expect(omitUnusedArgs(args, info)).toEqual([
+      'arg1',
+      'arg2',
+      'arg3',
+      'arg4',
+      'arg5',
+      'arg6',
+      undefined,
+      undefined,
+    ])
+  })
+
+  test('should handle args array shorter than 6 elements', () => {
+    const args = ['arg1', 'arg2', 'arg3']
+
+    const info: ServerReferenceInfo = {
+      type: 'server-action',
+      usedArgs: [true, false, true, false, false, false],
+      hasRestArgs: false,
+    }
+
+    expect(omitUnusedArgs(args, info)).toEqual(['arg1', undefined, 'arg3'])
+  })
+
+  test('should handle empty args array', () => {
+    const args: unknown[] = []
+
+    const info: ServerReferenceInfo = {
+      type: 'server-action',
+      usedArgs: [false, false, false, false, false, false],
+      hasRestArgs: false,
+    }
+
+    expect(omitUnusedArgs(args, info)).toEqual(args)
+  })
+})

packages/next/src/client/components/router-reducer/reducers/server-reference-info.ts

@@ -0,0 +1,78 @@
+export interface ServerReferenceInfo {
+  type: 'server-action' | 'use-cache'
+  usedArgs: [boolean, boolean, boolean, boolean, boolean, boolean]
+  hasRestArgs: boolean
+}
+
+/**
+ * Extracts info about the server reference for the given server reference ID by
+ * parsing the first byte of the hex-encoded ID.
+ *
+ * ```
+ * Bit positions: [7]      [6] [5] [4] [3] [2] [1]  [0]
+ * Bits:          typeBit  argMask                  restArgs
+ * ```
+ *
+ * If the `typeBit` is `1` the server reference represents a `"use cache"`
+ * function, otherwise a server action.
+ *
+ * The `argMask` encodes whether the function uses the argument at the
+ * respective position.
+ *
+ * The `restArgs` bit indicates whether the function uses a rest parameter. It's
+ * also set to 1 if the function has more than 6 args.
+ *
+ * @param id hex-encoded server reference ID
+ */
+export function extractInfoFromServerReferenceId(
+  id: string
+): ServerReferenceInfo {
+  const infoByte = parseInt(id.slice(0, 2), 16)
+  const typeBit = (infoByte >> 7) & 0x1
+  const argMask = (infoByte >> 1) & 0x3f
+  const restArgs = infoByte & 0x1
+  const usedArgs = Array(6)
+
+  for (let index = 0; index < 6; index++) {
+    const bitPosition = 5 - index
+    const bit = (argMask >> bitPosition) & 0x1
+    usedArgs[index] = bit === 1
+  }
+
+  return {
+    type: typeBit === 1 ? 'use-cache' : 'server-action',
+    usedArgs: usedArgs as [
+      boolean,
+      boolean,
+      boolean,
+      boolean,
+      boolean,
+      boolean,
+    ],
+    hasRestArgs: restArgs === 1,
+  }
+}
+
+/**
+ * Creates a sparse array containing only the used arguments based on the
+ * provided action info.
+ */
+export function omitUnusedArgs(
+  args: unknown[],
+  info: ServerReferenceInfo
+): unknown[] {
+  const filteredArgs = new Array(args.length)
+
+  for (let index = 0; index < args.length; index++) {
+    if (
+      (index < 6 && info.usedArgs[index]) ||
+      // This assumes that the server reference info byte has the restArgs bit
+      // set to 1 if there are more than 6 args.
+      (index >= 6 && info.hasRestArgs)
+    ) {
+      filteredArgs[index] = args[index]
+    }
+  }
+
+  return filteredArgs
+}

test/e2e/app-dir/use-cache/app/use-action-state/cached.ts

@@ -0,0 +1,7 @@
+'use cache'
+
+export async function getRandomValue() {
+  const v = Math.random()
+  console.log(v)
+  return v
+}

test/e2e/app-dir/use-cache/app/use-action-state/page.tsx

@@ -0,0 +1,15 @@
+'use client'
+
+import { useActionState } from 'react'
+import { getRandomValue } from './cached'
+
+export default function Page() {
+  const [result, formAction, isPending] = useActionState(getRandomValue, -1)
+
+  return (
+    <form action={formAction}>
+      <button id="submit-button">Submit</button>
+      <p>{isPending ? 'loading...' : result}</p>
+    </form>
+  )
+}

test/e2e/app-dir/use-cache/use-cache.test.ts

@@ -373,4 +373,24 @@ describe('use-cache', () => {
 
     expect(await browser.elementByCss('#random').text()).toBe(initialValue)
   })
+
+  it('works with useActionState if previousState parameter is not used in "use cache" function', async () => {
+    const browser = await next.browser('/use-action-state')
+
+    let value = await browser.elementByCss('p').text()
+    expect(value).toBe('-1')
+
+    await browser.elementByCss('button').click()
+
+    await retry(async () => {
+      value = await browser.elementByCss('p').text()
+      expect(value).toMatch(/\d\.\d+/)
+    })
+
+    await browser.elementByCss('button').click()
+
+    await retry(async () => {
+      expect(await browser.elementByCss('p').text()).toBe(value)
+    })
+  })
 })