Custom storage implementation(s)

[illusionalsagacity]

https://jotai.org/docs/utilities/storage#atomwithstorage

I see two obvious paths for this:

1. A zero-cost binding just with a record type and functions in it for both sync and async storage
2. A module function approach

Record Type

Here's an example binding for it

type unsubscribe = unit => unit
type subscribe<'value> = 'value => unit

type syncStorageConfig<'value> = {
  getItem: (string, 'value) => 'value,
  setItem: (string, 'value) => unit,
  removeItem: string => unit,
  subscribe?: (string, subscribe<'value>, 'value) => unsubscribe,
}

@module("jotai/utils")
external makeWithSyncConfig: (
  string,
  'value,
  syncStorageConfig<'value>,
) => Atom.t<'value, Atom.Actions.t<'value>, [Atom.Tags.r | Atom.Tags.w]> = "atomWithStorage"

The thing I dislike about this approach is it removes the ability to use a variant for the key without conversion. Or have the key be more complex like so:

type key = Todo({id: string}) | User({id: string})

Module Function

2 seems nicer for larger applications where you want to share a Storage implementation across several atoms, but it is less approachable / diverges from the JavaScript API.

This could look like:

type unsubscribe = unit => unit
type subscribe<'value> = 'value => unit

module type SyncStorageConfiguration = {
  type key // this doesn't quite work as it makes the key type private for any implementation
  type value

  let getItem: (key, value) => value
  let setItem: (key, value) => unit
  let removeItem: key => unit
  let subscribe: option<(key, subscribe<value>, value) => unsubscribe>
}

I like how this makes the storage implementation separate from "how I configure Jotai" but I can see this being considered overly opinionated.

@Fattafatta let me know which approach you like (or even if you have other ideas) and I'd be happy to contribute it

[Fattafatta]

Hi @illusionalsagacity, thanks for your suggestions. I'll have a look into it. The first one seems pretty straight forward. So I would definitely go with that. But that's not an either/or decision. We can simply add both. Perhaps using a Functor could solve the problem with the private key type.

[Fattafatta]

I experimented a little with this and came up with two solutions.

Generic 'key type

The first one is actually almost identical to your first suggestion. I just replaced the string type with a generic 'key type. Now it works for any kind of key:

type unsubscribe = unit => unit
type syncStorageConfig<'key, 'value> = {
  getItem: ('key, 'value) => 'value,
  setItem: ('key, 'value) => unit,
  removeItem: 'key => unit,
  subscribe: option<('key, 'value => unit, 'value) => unsubscribe>,
}

@module("jotai/utils")
external makeWithStorage: (
  'key,
  'value,
  syncStorageConfig<'key, 'value>,
) => Atom.t<'value, Atom.Actions.set<'value>, [Atom.Tags.r | Atom.Tags.w]> = "atomWithStorage"

Here is a (very simplistic) example:

type key = Add(int) | Mult(int)
let a = Utils.AtomWithStorage.makeWithStorage(
  Add(0),
  1,
  {
    getItem: (key, initialVal) => {
      switch key {
      | Add(i) => i + initialVal
      | Mult(i) => i * initialVal
      }
    },
    setItem: (_, _) => (),
    removeItem: _ => (),
    subscribe: None,
  },
)

The storage implementation can easily be shared by just storing the config.

With a Functor

This is a little more complicated, but instead of sharing the config record, you can create a new module with a given storage implementation, and then simply reuse the module.

module type Config = {
  type key
  type value

  let getItem: (key, value) => value
  let setItem: (key, value) => unit
  let removeItem: key => unit
  let subscribe: option<(key, value => unit, value) => unsubscribe>
}

module Make = (Config: Config) => {
  let _makeConfig = () => {
    getItem: Config.getItem,
    setItem: Config.setItem,
    removeItem: Config.removeItem,
    subscribe: Config.subscribe,
  }

  @module("jotai/utils")
  external _make: (
    Config.key,
    Config.value,
    syncStorageConfig<Config.key, Config.value>,
  ) => Atom.t<Config.value, Atom.Actions.set<Config.value>, [Atom.Tags.r | Atom.Tags.w]> =
    "atomWithStorage"

  let make = (key, initialValue) => _make(key, initialValue, _makeConfig())
}

Example:

module FakeStorage = {
  type key = Add(int) | Mult(int)
  type value = int

  let getItem = (key, initialVal) => {
    switch key {
    | Add(i) => i + initialVal
    | Mult(i) => i * initialVal
    }
  }
  let setItem = (_, _) => ()
  let removeItem = _ => ()
  let subscribe = Some((_, _, _) => () => ())
}
module FS = Utils.AtomWithStorage.Make(FakeStorage)

let a = FS.make(Add(0), 1)

@illusionalsagacity which variant would you prefer? I think, the first one is much simpler, but the second one is more elegant with regards to reusabilty.

[illusionalsagacity]

@illusionalsagacity which variant would you prefer? I think, the first one is much simpler, but the second one is more elegant with regards to reusabilty.

I prefer the functor since it fits nicely with SafeLocalStorage / SafeSessionStorage modules we're using to catch exceptions from accessing storage in private browsing--but I think it'd be nice to also expose the external make binding for those with a simpler use-case.

Would the plan for async storage be a MakeAsync functor? Or would a Sync / Async submodule make sense to you?