o1js/src/lib/provable/option.ts

TypeScript framework for zk-SNARKs and zkApps

import { InferValue } from '../../bindings/lib/provable-generic.js';
import { emptyValue } from '../proof-system/zkprogram.js';
import { Provable } from './provable.js';
import { InferProvable, Struct } from './types/struct.js';
import { provable, ProvableInferPureFrom } from './types/provable-derivers.js';
import { Bool } from './wrapped.js';

export { Option, OptionOrValue };

type Option<T, V = any> = { isSome: Bool; value: T } & {
  assertSome(message?: string): T;
  orElse(defaultValue: T | V): T;
};

type OptionOrValue<T, V> =
  | { isSome: boolean | Bool; value: T | V }
  | T
  | V
  | undefined;

/**
 * Define an optional version of a provable type.
 *
 * @example
 * ```ts
 * class OptionUInt64 extends Option(UInt64) {}
 *
 * // create an optional UInt64
 * let some = OptionUInt64.from(5n);
 * let none = OptionUInt64.none();
 *
 * // get back a UInt64
 * let five: UInt64 = some.assertSome('must have a value');
 * let zero: UInt64 = none.orElse(0n); // specify a default value
 * ```
 */
function Option<A extends Provable<any, any>>(
  type: A
): ProvableInferPureFrom<
  A,
  Option<InferProvable<A>, InferValue<A>>,
  InferValue<A> | undefined
> &
  (new (option: { isSome: Bool; value: InferProvable<A> }) => Option<
    InferProvable<A>,
    InferValue<A>
  >) & {
    fromValue(
      value:
        | { isSome: boolean | Bool; value: InferProvable<A> | InferValue<A> }
        | InferProvable<A>
        | InferValue<A>
        | undefined
    ): Option<InferProvable<A>, InferValue<A>>;
    from(
      value?: InferProvable<A> | InferValue<A>
    ): Option<InferProvable<A>, InferValue<A>>;
    none(): Option<InferProvable<A>, InferValue<A>>;
  } {
  type T = InferProvable<A>;
  type V = InferValue<A>;
  let strictType: Provable<T, V> = type;

  // construct a provable with a JS type of `T | undefined`
  const PlainOption: Provable<
    { isSome: Bool; value: T },
    { isSome: boolean; value: V }
  > = provable({ isSome: Bool, value: strictType });

  const RawOption = {
    ...PlainOption,

    toValue({ isSome, value }: { isSome: Bool; value: T }) {
      return isSome.toBoolean() ? strictType.toValue(value) : undefined;
    },

    fromValue(value: OptionOrValue<T, V>) {
      if (value === undefined)
        return { isSome: Bool(false), value: emptyValue(strictType) };
      // TODO: this isn't 100% robust. We would need recursive type validation on any nested objects to make it work
      if (typeof value === 'object' && 'isSome' in value)
        return PlainOption.fromValue(value as any); // type-cast here is ok, matches implementation
      return { isSome: Bool(true), value: strictType.fromValue(value) };
    },
  };

  const Super = Struct(RawOption);
  return class Option_ extends Super {
    orElse(defaultValue: T | V): T {
      return Provable.if(
        this.isSome,
        strictType,
        this.value,
        strictType.fromValue(defaultValue)
      );
    }

    assertSome(message?: string): T {
      this.isSome.assertTrue(message);
      return this.value;
    }

    static from(value?: V | T) {
      return value === undefined
        ? new Option_({ isSome: Bool(false), value: emptyValue(strictType) })
        : new Option_({
            isSome: Bool(true),
            value: strictType.fromValue(value),
          });
    }
    static none() {
      return Option_.from(undefined);
    }

    static fromFields(fields: any[], aux?: any): Option_ {
      return new Option_(Super.fromFields(fields, aux));
    }
    static fromValue(value: OptionOrValue<T, V>) {
      return new Option_(Super.fromValue(value));
    }
  };
}