o1-labs / o1js

TypeScript framework for zk-SNARKs and zkApps
https://docs.minaprotocol.com/en/zkapps/how-to-write-a-zkapp
Apache License 2.0
479 stars 107 forks source link

union types #1558

Open harrysolovay opened 3 months ago

harrysolovay commented 3 months ago

Developers will want to dispatch different types of actions and reduce those actions based on type. Any suggestions as to how one might represent tagged unions in reducer action types?

I took a (not-yet-functioning) laymen's shot at implementing an Enum utility.

Enum.ts ```ts import { Field, FlexibleProvable, Provable, Struct, UInt64 } from "o1js" const tagI = Symbol() const tag = Symbol() export function Enum(variants: Variants) { const variantEntries = Object.entries(variants) const variantTagIds = Object.fromEntries(variantEntries.map(([tag], i) => [tag, i])) return class Variant extends Struct({ [tagI]: UInt64, ...variants }) { constructor(readonly tag: Tag, value: InstanceType) { super({ [tagI]: new UInt64(variantTagIds[tag]!), [tag]: value, ...Object.fromEntries( variantEntries.filter(([tag_]) => tag !== tag_).map(([tag_]) => [tag_, Field.empty()]), ), }) } match>(resultType: A, arms: MatchArms): T { return Provable.switch( variantEntries.map((_0, i) => (this[tagI as keyof this] as UInt64).equals(new UInt64(i))), resultType, variantEntries.map(([tag_]) => arms[tag_]!(this[this[tag as keyof this] as keyof this]! as never) ), ) } } } export type VariantsConstraint = { [tag: string]: new(...args: any) => any } export type MatchArms = { [Tag in keyof Variants]: (value: InstanceType) => Return } ```

My hope was that this would enable me to define my reducer's actionType like so.

class Person extends Struct({
  age: UInt8,
  centenarian: Bool,
}) {}

class Dog extends Struct({
  age: UInt8,
  denarian: Bool,
}) {}

const Animal = Enum({ Person, Dog })

export class EnumTest extends SmartContract {
  reducer = Reducer({ actionType: Animal })
  // ...
}

Next, I would hope to initialize and send txs that contain variants of Animal / dispatch those variants as actions.

@method
async foo() {
  const doggo = new Animal("Dog", {
    age: new UInt8(4),
    denarian: new Bool(false),
  })
  this.reducer.dispatch(doggo)
}

Finally, I'd like to use match in my actions reducer to process specific variants depending on their tag.

@method
async rollup() {
  const state = this.state.getAndRequireEquals()
  const pending = this.reducer.getActions({ fromActionState: state.rollup })
  const reduced = this.reducer.reduce(
    pending,
    AnimalCounts,
    (state, animal) => {
      return animal.match(AnimalCounts, {
        Dog() {
          return new AnimalCounts({
            dogs: state.dogs.add(1),
            people: state.people,
          })
        },
        Person() {
          return new AnimalCounts({
            dogs: state.dogs,
            people: state.people.add(1),
          })
        },
      })
    },
    { state, actionState: state.rollup },
  )
  this.state.set(reduced.state)
}

Note: the first argument of the match method is the type we wish to return from the match arms / variant-specific handlers. This type is passed––under the hood––to a Provable.switch call.

mitschabaude commented 3 months ago

I think this is a very pretty way of handling enums, would be so nice to improve the experience.

I'm just a bit wary of promoting the pattern of defining arbitrary functions in conditional branches.

why? because it suggests that just one of these branches will execute (when in reality, all of them execute).

so I'm worried that users will put "conditional constraints" in those branches, and then would be confused why these constraints fail even if the branch was not supposed to be active.

wdyt @harrysolovay?

mitschabaude commented 3 months ago

consider this unsatisfiable circuit which your proposed API makes easy to write:

@method myMethod(animal: Animal, isDog: Bool) {
  animal.match(AnimalCounts, {
    Dog() {
      isDog.assertTrue();
    },
    Person() {
      isDog.assertFalse();
    },
  });
}
harrysolovay commented 3 months ago

when in reality, all of them execute

Is this a problem? If we find a way to force the irrelevant branches to pass, could we just pretend the irrelevant branches don't exist, or do they have implications on performance and scalability?

mitschabaude commented 3 months ago

If we find a way to force the irrelevant branches to pass, could we just pretend the irrelevant branches don't exist, or do they have implications on performance and scalability?

I'd love to implement proper branching and have a plan for it, but it's not exactly a simple problem that you quickly find a way for. The constraints a circuit uses are fundamentally static, so if in the example above, the circuit really does contain both statements assert(dog = true) and assert(dog = false), then this is fundamentally not a circuit you can prove, no way to "make it pass"

mitschabaude commented 3 months ago

and yes, any branching solution that adds the ability to disable branches creates overhead constraints (so it should be a separate API)

harrysolovay commented 3 months ago

I've been thinking a lot about this issue. And thank you @mitschabaude for your patience: I don't yet understand the underlaying limitation, but I'm sure it's no small feat to circumvent that limitation. That being said, I believe a lack of polymorphic types will severely limit the kinds of programs that can be expressed with o1js. Do you see any stopgap solution to emit different types of action and apply different reducer callbacks based on type?

mitschabaude commented 3 months ago

Hey @harrysolovay originally I thought a possible solution would be to make match more similar to Provable.if: just selecting between values not functions.

But now I realize functions are essential in order to easily get the content of the matched variant.

So I'd be in favor of trying your idea and heavily documenting how not to use it. Because if you know how to use it, it's a very nice API

mitschabaude commented 3 months ago

Another interesting alternative proposed by @nholland94 is to have a match() function that can directly apply to a native TS union like Dog | Person. Quoting @nholland94:

For ergonomics: instead of having an Enum(...) constructor that assigns local discriminates to the union, you could use global discriminators by having each taggable class (Tagged) give an interface to an embedded discriminator which is computed by hashing the name of the type. A match function, accepting a list of Tagged class instances to match against, and functions to dispatch on when matched, would do a quick out-of-circuit check that all the discriminators are unique when the circuit is compiled. So long as a consistent and portable hash function is used, should mostly just work, and doesn't require JS devs to define a custom Enum type. The Enum type route is also interesting, and you could do some smart things with it by generating nice helpers based on the enum's inhabitants, but then you get into all that weird complex TypeScript stuff to deal with dynamically generated objects, which people can get confused by.

In that world, I would have a match function like:

const myValue: Taggable = ...;
match(myValue, [
  [X, (x: X) => ...],
  [Y, (y: Y) => ...],
])

In this, the function body for a match case doesn't pattern match out all of the subelements, like in ocaml or haskell, but rather, casts the outer taggable type to a better understood inner type. I think you can represent this in typescript, not 100% sure. The alternatives of having to read out the members to pass into a function call will either end up being really cumbersome type-wise, or will have to introduce type unsafety by casting to and from any.

harrysolovay commented 3 months ago

have a match() function that can directly apply to a native TS union like Dog | Person

Event better!

then you get into all that weird complex TypeScript stuff to deal with dynamically generated objects

I'm not sure I follow. Can you please say this in a different way?


An alternative API suggestion:

declare const animal: Dog | Person

match(animal, _ => _
  .when(Dog, (value) => {
    value // `Dog`
  })
  .when(Person, (value) => {
    value // `Person`
  })
)

Perhaps the 1st param of when can be typed such that one can supply a guard as well. This would enable more fine-grained progressive narrowing of the matched type.

function isDog(value: Dog | Person): value is Dog {
  return value instanceof Dog
}
function isPerson(value: Dog | Person): value is Person {
  return value instanceof Person
}

match(animal, _ => _
  .when(isDog, (value) => {
    value // `Dog`
  })
  .when(isPerson, (value) => {
    value // `Person`
  })
)

Additionally, the returned "matcher" chain can be type-checked for exhaustiveness.

declare const animal: Dog | Person | SlowLoris

match(animal, _ => _
  .when(Dog, (value) => {
    value // `Dog`
  })
  .when(Person, (value) => {
    value // `Person`
  })
)

^ this could produce a type error, given that the SlowLoris case is not handled. This could be fixed either by adding another when to the chain, or by using an else.

match(animal, _ => _
  .when(Dog, (value) => {
    value // `Dog`
  })
  .when(Person, (value) => {
    value // `Person`
  })
  .else((value) => {
    value // `SlowLoris`
  })
)
mitschabaude commented 3 months ago

love the .when() syntax

harrysolovay commented 3 months ago

Shoutout to @tjjfvi who implemented this design for another virtual language back in the day (see here).