feat(pic): support canister snapshots

docs/src/content/docs/guides/canister-declarations.mdx

@@ -1,7 +1,7 @@
 ---
 title: Canister declarations
 sidebar:
-  order: 7
+  order: 8
 ---
 
 PicJS creates actors from a canister's Candid interface: an `idlFactory` that encodes and decodes calls, and a `_SERVICE` type that describes them. Generate both from the canister's `.did` file with [`@icp-sdk/bindgen`](https://js.icp.build/bindgen/latest/):

docs/src/content/docs/guides/canister-snapshots.mdx

@@ -0,0 +1,58 @@
+---
+title: Canister snapshots
+sidebar:
+  order: 7
+---
+
+A [canister snapshot](https://docs.internetcomputer.org/guides/canister-management/snapshots/) captures a canister's state: its Wasm module, heap and stable memory, and chunk store. All snapshot methods must be called by a controller of the canister.
+
+## Restoring state
+
+Take a snapshot, change the canister's state, and load the snapshot to go back, for example to roll back a failed upgrade:
+
+```ts
+const snapshot = await pic.takeCanisterSnapshot({ canisterId, sender });
+
+await pic.upgradeCanister({ canisterId, wasm, sender });
+// assert on the upgraded canister...
+
+await pic.loadCanisterSnapshot({
+  canisterId,
+  snapshotId: snapshot.id,
+  sender,
+});
+```
+
+Snapshots can also be listed with `listCanisterSnapshots` and deleted with `deleteCanisterSnapshot`. Pass `replaceSnapshot` to `takeCanisterSnapshot` to replace an existing snapshot instead of adding one.
+
+## Testing against real state
+
+Snapshots are not tied to a network, so a snapshot downloaded from mainnet or a local network can be uploaded into PocketIC to test against real state. Download it with [icp-cli](https://cli.internetcomputer.org):
+
+```shell
+icp canister snapshot create <canister> -n ic
+icp canister snapshot download <canister> <snapshot-id> -n ic -o ./snapshot
+```
+
+Then upload it to a canister in PocketIC and load it:
+
+```ts
+const canisterId = await pic.createCanister({ sender });
+
+const snapshotId = await pic.uploadCanisterSnapshot({
+  canisterId,
+  snapshotDir: path.resolve('./snapshot'),
+  sender,
+});
+await pic.loadCanisterSnapshot({ canisterId, snapshotId, sender });
+```
+
+`downloadCanisterSnapshot` writes a snapshot in the same format, which `icp canister snapshot upload` accepts.
+
+`snapshotDir` is a path on the machine running the PocketIC server, and should be absolute.
+
+The snapshot carries state that may depend on where it was taken:
+
+- The canister ID changes, so state that stores the canister's own ID refers to the original canister.
+- Other canisters the state refers to may not exist on PocketIC under the same IDs.
+- Public keys derived from threshold keys (ECDSA, Schnorr, vetKD) differ between networks.

docs/src/content/docs/guides/more-examples.mdx

@@ -2,7 +2,7 @@
 title: More Examples
 next: false
 sidebar:
-  order: 8
+  order: 9
 ---
 
 All examples are written in [TypeScript](https://www.typescriptlang.org/) with [Jest](https://jestjs.io/) as the test runner,

packages/pic/src/management-canister.ts

@@ -441,3 +441,105 @@ export function decodeFetchCanisterLogsResponse(
 
   return payload;
 }
+
+const Snapshot = IDL.Record({
+  id: IDL.Vec(IDL.Nat8),
+  taken_at_timestamp: IDL.Nat64,
+  total_size: IDL.Nat64,
+});
+
+export interface Snapshot {
+  id: Uint8Array;
+  taken_at_timestamp: bigint;
+  total_size: bigint;
+}
+
+const TakeCanisterSnapshotRequest = IDL.Record({
+  canister_id: IDL.Principal,
+  replace_snapshot: IDL.Opt(IDL.Vec(IDL.Nat8)),
+  uninstall_code: IDL.Opt(IDL.Bool),
+  sender_canister_version: IDL.Opt(IDL.Nat64),
+});
+
+export interface TakeCanisterSnapshotRequest {
+  canister_id: Principal;
+  replace_snapshot: [] | [Uint8Array];
+  uninstall_code: [] | [boolean];
+  sender_canister_version: [] | [bigint];
+}
+
+export function encodeTakeCanisterSnapshotRequest(
+  arg: TakeCanisterSnapshotRequest,
+): Uint8Array {
+  return new Uint8Array(IDL.encode([TakeCanisterSnapshotRequest], [arg]));
+}
+
+export function decodeTakeCanisterSnapshotResponse(arg: Uint8Array): Snapshot {
+  const payload = decodeCandid<Snapshot>([Snapshot], arg);
+
+  if (isNil(payload)) {
+    throw new Error('Failed to decode TakeCanisterSnapshotResponse');
+  }
+
+  return payload;
+}
+
+const LoadCanisterSnapshotRequest = IDL.Record({
+  canister_id: IDL.Principal,
+  snapshot_id: IDL.Vec(IDL.Nat8),
+  sender_canister_version: IDL.Opt(IDL.Nat64),
+});
+
+export interface LoadCanisterSnapshotRequest {
+  canister_id: Principal;
+  snapshot_id: Uint8Array;
+  sender_canister_version: [] | [bigint];
+}
+
+export function encodeLoadCanisterSnapshotRequest(
+  arg: LoadCanisterSnapshotRequest,
+): Uint8Array {
+  return new Uint8Array(IDL.encode([LoadCanisterSnapshotRequest], [arg]));
+}
+
+const ListCanisterSnapshotsRequest = IDL.Record({
+  canister_id: IDL.Principal,
+});
+
+export interface ListCanisterSnapshotsRequest {
+  canister_id: Principal;
+}
+
+export function encodeListCanisterSnapshotsRequest(
+  arg: ListCanisterSnapshotsRequest,
+): Uint8Array {
+  return new Uint8Array(IDL.encode([ListCanisterSnapshotsRequest], [arg]));
+}
+
+export function decodeListCanisterSnapshotsResponse(
+  arg: Uint8Array,
+): Snapshot[] {
+  const payload = decodeCandid<Snapshot[]>([IDL.Vec(Snapshot)], arg);
+
+  if (isNil(payload)) {
+    throw new Error('Failed to decode ListCanisterSnapshotsResponse');
+  }
+
+  return payload;
+}
+
+const DeleteCanisterSnapshotRequest = IDL.Record({
+  canister_id: IDL.Principal,
+  snapshot_id: IDL.Vec(IDL.Nat8),
+});
+
+export interface DeleteCanisterSnapshotRequest {
+  canister_id: Principal;
+  snapshot_id: Uint8Array;
+}
+
+export function encodeDeleteCanisterSnapshotRequest(
+  arg: DeleteCanisterSnapshotRequest,
+): Uint8Array {
+  return new Uint8Array(IDL.encode([DeleteCanisterSnapshotRequest], [arg]));
+}

packages/pic/src/pocket-ic-client-types.ts

@@ -411,6 +411,66 @@ export function encodeCreateInstanceRequest(
 
 //#endregion CreateInstance
 
+//#region CanisterSnapshotTransfer
+
+export interface CanisterSnapshotDownloadRequest {
+  sender: Principal;
+  canisterId: Principal;
+  snapshotId: Uint8Array;
+  snapshotDir: string;
+}
+
+export interface EncodedCanisterSnapshotDownloadRequest {
+  sender: { principal_id: string };
+  canister_id: { canister_id: string };
+  snapshot_id: string;
+  snapshot_dir: string;
+}
+
+export function encodeCanisterSnapshotDownloadRequest(
+  req: CanisterSnapshotDownloadRequest,
+): EncodedCanisterSnapshotDownloadRequest {
+  return {
+    sender: { principal_id: base64EncodePrincipal(req.sender) },
+    canister_id: { canister_id: base64EncodePrincipal(req.canisterId) },
+    snapshot_id: base64Encode(req.snapshotId),
+    snapshot_dir: req.snapshotDir,
+  };
+}
+
+export interface CanisterSnapshotUploadRequest {
+  sender: Principal;
+  canisterId: Principal;
+  replaceSnapshot?: Uint8Array;
+  snapshotDir: string;
+}
+
+export interface EncodedCanisterSnapshotUploadRequest {
+  sender: { principal_id: string };
+  canister_id: { canister_id: string };
+  replace_snapshot: { snapshot_id: string } | null;
+  snapshot_dir: string;
+}
+
+export function encodeCanisterSnapshotUploadRequest(
+  req: CanisterSnapshotUploadRequest,
+): EncodedCanisterSnapshotUploadRequest {
+  return {
+    sender: { principal_id: base64EncodePrincipal(req.sender) },
+    canister_id: { canister_id: base64EncodePrincipal(req.canisterId) },
+    replace_snapshot: req.replaceSnapshot
+      ? { snapshot_id: base64Encode(req.replaceSnapshot) }
+      : null,
+    snapshot_dir: req.snapshotDir,
+  };
+}
+
+export interface EncodedCanisterSnapshotUploadResponse {
+  snapshot_id: string;
+}
+
+//#endregion CanisterSnapshotTransfer
+
 //#region GetPubKey
 
 export interface GetPubKeyRequest {

packages/pic/src/pocket-ic-client.ts

@@ -2,6 +2,13 @@ import { JSONParse } from 'json-with-bigint';
 import { Http2Client } from './http2-client';
 import { ServerRequestTimeoutError } from './error';
 import {
+  CanisterSnapshotDownloadRequest,
+  CanisterSnapshotUploadRequest,
+  EncodedCanisterSnapshotDownloadRequest,
+  EncodedCanisterSnapshotUploadRequest,
+  EncodedCanisterSnapshotUploadResponse,
+  encodeCanisterSnapshotDownloadRequest,
+  encodeCanisterSnapshotUploadRequest,
   EncodedAddCyclesRequest,
   EncodedAddCyclesResponse,
   EncodedCanisterCallRequest,
@@ -89,7 +96,7 @@ import {
   EncodedHttpGatewayRequest,
   EncodedHttpGatewayResponse,
 } from './pocket-ic-client-types';
-import { base64DecodePrincipal, isNil } from './util';
+import { base64Decode, base64DecodePrincipal, isNil } from './util';
 import { Principal } from '@icp-sdk/core/principal';
 
 const PROCESSING_TIME_VALUE_MS = 30_000;
@@ -258,6 +265,33 @@ export class PocketIcClient {
     return decodeAddCyclesResponse(res);
   }
 
+  public async canisterSnapshotDownload(
+    req: CanisterSnapshotDownloadRequest,
+  ): Promise<void> {
+    this.assertInstanceNotDeleted();
+
+    await this.post<EncodedCanisterSnapshotDownloadRequest, {}>(
+      '/update/canister_snapshot_download',
+      encodeCanisterSnapshotDownloadRequest(req),
+    );
+  }
+
+  public async canisterSnapshotUpload(
+    req: CanisterSnapshotUploadRequest,
+  ): Promise<Uint8Array> {
+    this.assertInstanceNotDeleted();
+
+    const res = await this.post<
+      EncodedCanisterSnapshotUploadRequest,
+      EncodedCanisterSnapshotUploadResponse
+    >(
+      '/update/canister_snapshot_upload',
+      encodeCanisterSnapshotUploadRequest(req),
+    );
+
+    return base64Decode(res.snapshot_id);
+  }
+
   public async uploadBlob(req: UploadBlobRequest): Promise<UploadBlobResponse> {
     this.assertInstanceNotDeleted();