Problem with polymorphic bindings

I’m currently trying to make my bindings for jotai (state management for react) a little more robust. But it’s getting frustratingly complicated.

Jotai’s main building block is an atom. Which is basically a polymorphic function that returns a container that holds some state. The container can hold basically any type. Container creation is very flexible. It accepts primitive types, functions, async function and many more. Whatever is passed to the function is getting resolved before it is put in the container.
So for example if you pass in a function, the function gets called and the return value is put in the container.
If you pass in an async function, it is called and the promise resolved before the resolved value is put in the container.

(I’m simplifying here, the actual atom is a bit more complex, but this should be sufficient to clarify the problem)

My current solution is that I create multiple functions that bind to the same function from jotai:

// Simple atom
external atom: 'value => 'value = "atom"
let fromInt = atom(1)
let fromList = atom([1,2 ,3])

// Atom from function
external atomFunc: (unit =>'value) => 'value = "atom"
let fromFunc = atomFunc(() => 1)

// Atom from async function
external atomAsync: (unit => promise<'value>) => 'value = "atom"
let fromPromise = atomAsync(async () => 1)

Code: Playground
Repo with actual bindings: rescript-jotai

This is generally working in the sense that I can support almost any scenario implemented by jotai. The (really annoying) problem is, that it is not type save.

For example this will compile but cause a runtime error:

// this will cause trouble because the actual return type is "int"
let willFail = atomFunc(async () => 1)

And the atom function with its generic definition ('value => 'value) is even worse. It literally accepts anything and will mess up the return type.

Is there any way to somehow constrain those bindings in some way? So for example would it be possible to create a type like “accept any input that is not a function” or “accept only synchronous functions”.

As far as I understand, this is not possible, but perhaps you have some ideas for suitable workarounds.
Any help is highly appreciated!

1 Like

Interesting question…

What do you think of a solution like this (playground link)?

module Atom = {
  type t<'a>

  // Literally anything will work!
  external atom_magic: 'a => t<'a> = "atom"

  // Type safe examples...put as many of these as you will need.
  external i: int => t<int> = "atom"
  external f: float => t<float> = "atom"
  external s: string => t<string> = "atom"

  // You may notice that you can pass the async functions to thunk 
  // as well...see below.
  external thunk: (unit => 'a) => t<'a> = "atom"

  // Note the type here:  you say that by some external magic,
  // jotai knows to get the actual value out of the promise
  // and put that value into the atom rather than the promised value, 
  // so that's what the type signature says is happening.
  // (This is unlike the thunk function above.)
  external promise: (unit => promise<'a>) => t<'a> = "atom"

  // Pretend there's a function to get the value of the atom.
  // ...there may be, I don't know, but we use this for illustration.
  external unwrap: t<'a> => 'a = "unwrap"
}

// Atom.i("won't work")->ignore // compiler error!

// okay!
let int_atom: Atom.t<int> = Atom.promise(async () => 1)
let y = Atom.unwrap(int_atom) + 1

So back to your problem of passing promises to the thunk function…

Here, we defined thunk as an external to “atom”. So, if you give it a thunk that returns a promise (which will typecheck just fine) jotai will (presumably) do it’s thing and stick an int in the atom. But you told the ReScript compiler that what you get back is promise, which isn’t correct. It compiles fine, but the behavior is modeled incorrectly.

But how then will you know you’ve made an error? (I.e., how will you know you should have used the promise function instead of the thunk function…) Well, when you try and use the value incorrectly downstream, or you annotate it improperly, the compiler will complain.

Check it out.


In this example, the types don’t match the way that I have annotated them, so I know I made a mistake and should have used Atom.promise instead.

// Error: This has type: 
// Atom.t<promise<int>> 
//   ...Somewhere wanted: Atom.t<int>
let wont_work: Atom.t<int> = Atom.thunk(async () => 1) 

In the following example, I try to use the unwrapped value as an int, but I got a type error because it is a promised int. What happened? Oh, I made a mistake above and should have used Atom.promise, instead of Atom.thunk.

// Error: This has type: promise<int> ...Somewhere wanted: int
let wont_work = Atom.unwrap(Atom.thunk(async () => 1)) + 1 

What about if you don’t have the unwrap function? It could be that that your library doesn’t give a way out of the Atom type. How will I know I’ve made the modeling error in that case?

Well, presumably you will be writing functions that consume those created atoms, so just restrict the type that those function parameters should be …

let f: Atom.t<int> => int = _x => {
  failwith("omitted")
}

// This won't work....
let y = {
  let x = Atom.thunk(async () => 1)
  // Can't use x here as it is the wrong type
  f(x) // Error: This has type: Atom.t<promise<int>> ...Somewhere wanted: Atom.t<int>
}

// This is okay!
let z = {
  let x = Atom.promise(async () => 1)
  f(x)
}

Now, it would still be a domain modeling error if you pass a promise-returning function into the thunk function. That’s because the jotai library your binding to doesn’t work that way. So if you tried to do something like this

let a = Atom.thunk(async () => 1)
let f: Atom.t<promise<int>> => unit = _a => failwith("omitted")
let y = f(a)

Well, that will compile fine, but it is still wrong! But from my understanding of your explanation of jotai, a promise can never be put into an atom, so presumably you wouldn’t write a function that takes a parameter of type Atom.t<promise<'a>> anyway.

3 Likes

@Ryan thanks a lot for the detailed answer. The unwrap function you describe is one of the core problems with the bindings. In jotai this is a hook (called useAtom). And it behaves very similar to your unwrap function. With one critical distinction, all the magic described above only resurfaces when using this function. Perhaps I should have added useAtom to the example to make it more clear.

When you give an atom with a promise to useAtom it will return the resolved promise, if you give it a function it returns the result and so on. So the function signature looks more like this:

let useAtom: Atom.t<'value> => 'unwrappedvalue = atom => some_jotai_magic(atom)

And I’m currently searching for ways to make the return type less generic. Using multiple functions for different types is the approach I’m using right now. But it has limitations. In many cases nothing is preventing me from using the wrong function. The result would be that the actual type returned does not match the signature.

Like in your example, both these atoms would return an int when used with useAtom:

let a1 = promise(async () => 1)
let a2 = thunk(async () => 1)

// returns an int as expected
let v1 = useAtom(a1)
// returns an int but the compiler expects a promise
let v2 = useAtom(a2)

But, as you pointed out, whenever the value is used downstream, the wrong type should cause problems. So for someone who understands jotai, this would not be much of a problem. This person would know that useAtom would never return a promise etc.
But for someone new, this is not so easy to understand. For example there are other functions (getters and setters), that actually return the promise (and not the value), but the would still return the result of the function call (and not the function itself). So things get “partially unwrapped”. For beginners, this gets confusing very fast.

Jotai is written in TypeScript and they managed to create the correct type signatures for all these situations. But since the rescript type system is quite different (and in most cases far superior to typescript in my opinion), I’m trying to replicate that without success.

I have to think about your other suggestion of annotating the atoms at creation time. This would make the problem more explicit. But nothing would prevent me from annotating with the wrong types and the problem would only get transparent at runtime.

1 Like

Thanks for the clarification about the useAtom and getters function function…I think I can see more of the issue now.

It seems that the useAtom will “unwrap” as much as it needs to to get the value…whether by function application, waiting on the promise, whatever. But the getter is a bit different (I will call it get here). Something like this:

external getInt: atom<int> => int = "get"
external getThunk: atom<unit => int> => int = "get"
external getAsync: atom<promise<int>> => promise<int> = "get"
external getAsyncThunk: atom<unit => promise<int>> => promise<int> = "get"

Yeah the problem is still that you can pass an async thing (aka promise returning thing), to the function that is supposed to be the non-async version. And it compiles fine because unit => promise<'a> satisfies unit => 'a (ie promise<'a> is 'a)…

Take 2 (a bit better)

I could see something like this (playground). The phantom type keeps you from mixing up values created by the Sync and Async module functions.

module Atom = {
  type sync = Sync
  type async = Async
  type t<'value, 'color>

  module Sync = {
    external int: int => t<int, sync> = "atom"
    external useInt: t<int, sync> => int = "useAtom"
    external getInt: t<int, sync> => int = "get"

    type thunk<'a> = unit => 'a
    external thunk: thunk<'a> => t<thunk<'a>, sync> = "atom"
    external useThunk: t<thunk<'a>, sync> => 'a = "useAtom"
    external getThunk: t<thunk<'a>, sync> => 'a = "get"
  }

  module Async = {
    external int: int => t<promise<int>, async> = "atom"
    external useInt: t<promise<int>, async> => int = "useAtom"
    external getInt: t<promise<int>, async> => promise<int> = "get"

    type thunk<'a> = unit => promise<'a>
    external thunk: thunk<'a> => t<thunk<'a>, async> = "atom"
    external useThunk: t<thunk<'a>, async> => 'a = "useAtom"
    external getThunk: t<thunk<'a>, async> => promise<'a> = "get"
  }
}

It works pretty good except for when you pass an async thing to one of the sync modules. E.g., I can still pass a promise returning thunk to the Atom.Sync.thunk.

let x: Atom.t<Atom.Sync.thunk<promise<int>>, Atom.sync> = Atom.Sync.thunk(async () => 1)
let y: promise<int> = Atom.Sync.useThunk(x)

Which is definitely wrong at runtime as the library will “unwrap” that and y should be int as you say. But at least there you can clearly see something weird has happened if you look at the types…promise mixed in with Sync…oh something bad must have happened. Of course, it is still on the user to notice that.


Edit: Actually, I just thought of a different approach.

Take 3 (pretty good, if you don’t mind the functors)

The main problem is that you sync/async matters to the rescript type system, but the useAtom and get functions on the jotai side will do their thing regardless. The main problem was that you want to make sure you don’t send a promise to somewhere that can’t handle it properly. But the problem was the old thunk function (eg (unit => 'a) => atom<unit => 'a>), would happily take a unit => promise<'a> thunk, then the useAtom functions would return types that didn’t reflect what jotai was doing.

So, just don’t write any polymorphic bindings that take 'a. Make them all work with concrete types. Obviously, this would be really tedious to do manually, and, what if the user wants to use custom types, or records, or something you the library author can’t know about? Well, how about using functors. You can generate some common cases, and then let the user do whatever else they need to.

playground

module type T = {
  type t
}

module Atom = {
  type t<'a>

  module MakeAtomExternals = (T: T) => {
    module Sync = {
      external atom: T.t => t<T.t> = "atom"
      external use: t<T.t> => T.t = "useAtom"
      external get: t<T.t> => T.t = "get"

      external atomThunk: (unit => T.t) => t<unit => T.t> = "atom"
      external useThunk: t<unit => T.t> => T.t = "useAtom"
      external getThunk: t<unit => T.t> => T.t = "get"
    }

    module Async = {
      external atom: promise<T.t> => t<promise<T.t>> = "atom"
      external use: t<promise<T.t>> => T.t = "useAtom"
      external get: t<promise<T.t>> => promise<T.t> = "get"

      external atomThunk: (unit => promise<T.t>) => t<unit => promise<T.t>> =
        "atom"
      external useThunk: t<unit => promise<T.t>> => T.t = "useAtom"
      external getThunk: t<unit => promise<T.t>> => promise<T.t> = "get"
    }
  }

  // Include as many basic types as you want.
  module Int = MakeAtomExternals({
    type t = int
  })

  module String = MakeAtomExternals({
    type t = string
  })

  module Float = MakeAtomExternals({
    type t = string
  })

  module IntArray = MakeAtomExternals({
    type t = array<int>
  })
}

// These work fine.

let x = [1, 2, 3]->Atom.IntArray.Sync.atom->Atom.IntArray.Sync.use
let y =
  [1, 2, 3]->Js.Promise.resolve->Atom.IntArray.Async.atom->Atom.IntArray.Async.use

let x' =
  Atom.String.Async.atomThunk(async () => "yo")->Atom.String.Async.useThunk
let y' = Atom.String.Sync.atomThunk(() => "yo")->Atom.String.Sync.useThunk

// This won't work
// let oops = Atom.String.Sync.atom("yo")->Atom.String.Async.use

// And finally the dreaded sending async stuff to sync....fails to compile!
// let oops = Atom.String.Sync.atomThunk(async () => "yo")->Atom.String.Sync.useThunk

// Library consumer can use the functors on their custom types, too.
module Person = {
  module T = {
    type t = {name: string, age: int}
  }

  include T

  module Atom = {
    include Atom.MakeAtomExternals(T)
  }
}

let person =
  {Person.name: "Ryan", age: 35}
  ->Js.Promise.resolve
  ->Person.Atom.Async.atom
  ->Person.Atom.Async.use

// won't work either
//let oops = Person.Atom.Sync.atomThunk(async () => {Person.name: "Ryan", age: 35})

See how mixing up the sync and async are now compile time errors?

Of course, I’m still not exactly sure how the getters and setters work, but maybe you can take the functor approach and get closer to what you need.

@Ryan thanks a lot for the amount of thought you put into this. It is highly appreciated.

Your “Take 2” is very close to the solution I have right now. Im using polymorphic variants to “store certain attributes” for the atoms (I call them tags). For example an atom can be readonly, or primitive. I also added a “sync” tag for a while, but I removed it because it made the type more complex without providing much benefit. It was helpful in some cases, but the major problems (like differentiating between a thunk and an async thunk) didn’t go away.

Also, defining allowed primitive types beforehand is not really working, since you can put basically anything you want into an atom (primitives, but also arrays, objects, records, or even other atoms, …). So it is impossible do define them all beforehand.

Your “Take 3” on the other hand would be a good way to take care of some problems. You could create any type you want with it and use it with the right useAtom. This would require the user to create many Functors, but it is only one extra step to be almost completely type safe.

With almost I mean that it would be impossible to use the wrong useAtom function, but it would still be possible to use a Functor to build the wrong atom:

module IntWrong = MakeAtomExternals({
  type t = unit => int
})
// this would compile but have the wrong type
let a = (_ => 1)->Atom.IntWrong.Sync.atom->Atom.IntWrong.Sync.use

But although this would be possible, it is more obvious that something is wrong here. I will have to check if this remains viable when I try to add all the other stuff (like getters, setters and trags, etc.)

It’s a pity that there is no type like: “anything but a function”. With something like that, almost all problems would disappear. (Since we can handle the sync vs. async case quite well already)

Thanks again for the help!

Good point about the IntWrong module example…I guess you would have to trust the user to read docs. Or, I wonder if you could tweak things a bit, or make a functor dedicated to “function-like” types. Something to think about at least.

Just one thing to mention here…is that the user won’t have to write the functors, your library could provide them, and the lib consumer would only have to use them. (Possibly, this is what you meant anyway…)

This sounds like it could be an interesting post on the OCaml forum…there are many people there who could explain why the type system doesn’t work that way, and it would apply to ReScript as well.

Sure, no problem!