Storing function as jotai atom?

Question

I am wondering if it's possible to store a function as a jotai atom, and if so, how. Simply replicating the standard pattern for creating and consuming jotai atom-states, which works fine when the value of the state is an array, does not work when the value of the state is a function.

"./jotaiStore.js"

import { atom } from "jotai";

function test(x) {return x + 4;}
let testa = atom(test);
export { testa }; 

./App.js

import { testa } from "./jotaiStore.js";
import { useAtom } from "jotai";

export default function App() {
  
  const [locTest, locSetTest] = useAtom(testa);
  console.log(locTest);
  console.log(locTest(1));

  return (
    <div className="App">
    </div>
  );
}

See https://codesandbox.io/s/thirsty-brook-n1ucjr?file=/src/App.js:24-493.

The resulting log for console.log(locTest) is:

function (a) {
        dependencies.add(a);
        var aState = a === atom ? getAtomState(version, a) : readAtomState(version, a);

        if (aState) {
          if ('e' in aState) {
            throw aState.e;
          }

          if ('p' in aState) {
            throw aState.p;
          }

          return aState.v;
        }

        if (hasInitialValue(a)) {
          return a.init;
        }

        throw new Error('no atom init');
      }4 

and console.log(locTest(1)) generates an error "locTest is not a function".

EDIT: I found a way to store a function as a jotai state (see below) but not one I can --- or know how to --- update. If someone else has a better updatable solution, I'd still be very interested.

Answer 1

Not sure if you still need this, but Daishi actually suggests wrapping the function in an object.

const fnAtom = atom({ fn: () => "your function" })

You can then create a derived readable and/or writable atom based on the above primitive atom.

const doWhateverFnAtom = atom(
  (get) => get(fnAtom),
  (get, set, newFn) => {
    set(fnAtom, { fn: newFn })
  }
)

Source: https://github.com/pmndrs/jotai/discussions/1746#discussioncomment-4872129