Allow worklet referencing in plugin

[tjzel]

Why

Currently, when handling worklet callbacks, the user has either to mark a function directly with worklet directive or define the worklet as an inline argument.

// This will work:
const styleFactory = () => {
  'worklet';
  return { backgroundColor: 'blue' };
};
const animatedStyle = useAnimatedStyle(styleFactory);

// This will work as well:
const animatedStyle = useAnimatedStyle(() => ({ backgroundColor: blue }));

// However this won't
const styleFactory = () => {
  return { backgroundColor: 'blue' };
};
const animatedStyle = useAnimatedStyle(styleFactory); // error - style factory is not a worklet!

This pull request allows the user to define a worklet outside of a hook argument to get it autoworkletized. Keep in mind that it still has some boundaries:

1. Worklet has to be defined in the same file. It cannot be imported.
2. Worklet has to be defined before it's used.
3. Worklet cannot be defined via an expression (for example, using an ? operator). This however could be adjusted in the future.

What

We are allowing the following constructions to be autoworkletized. useAnimatedStyle will be used as a example hook here, it will work with every bit of our API that facilitates autoworkletization.

Function declarations

function worklet() {
  // Some UI relevant code.
}

const animatedStyle = useAnimatedStyle(worklet);

Function expressions

const worklet = function() {
  // Some UI relevant code.
}

const animatedStyle = useAnimatedStyle(worklet);

Arrow function expression

const worklet = () => {
  // Some UI relevant code.
}

const animatedStyle = useAnimatedStyle(worklet);

Object methods

This is a special case, used for example by useAnimatedScrollHandler.

const handler = {
  onScroll: () => {
    // Some UI relevant code.
  }
}

useAnimatedScrollHandler(handler);

It doesn't matter if the method is actually an arrow function or function expression. Keep in mind that it actually has to be defined in place. As of now we don't support such deep cases as:

const onScroll = () => {
  // Some UI relevant code.
}

const handler = {
  onScroll
}

useAnimatedScrollHandler(handler);

How

Changes implemented here are quite simple. The only thing we do is search the scope of a given context (that's usually a function) for a referenced identifier. For example:

const animatedStyle = useAnimatedStyle(styleFactory);

The styleFactory identifier is referenced in the useAnimatedStyle call. We then search the scope of the function for a reference to styleFactory and look for variable declarations, variable assignments, and function declarations.

1. Variable declarations:

   let styleFactory = () => {
     // Some UI relevant code.
   }

2. Variable assignments:

   let styleFactory;
   // ...
   styleFactory = () => {
     // Some UI relevant code.
   }

3. Function declarations:

   function styleFactory() {
     // Some UI relevant code.
   }

If we find one of these, that can be autoworkletized, we do it. If there are multiple objects that can be autoworkletized, we pick the first one using the following rules:

1. Function declarations are picked first.
2. Variable assignments are picked second. In case of multiple variable assignments, we pick the last one.
3. Variable declarations are picked last.

Therefore in the following code:

let styleFactory = function foo() {
  // Some UI relevant code.
}

styleFactory = function bar() {
  // Some UI relevant code.
}

styleFactory = function baz() {
  // Some UI relevant code.
}

const animatedStyle = useAnimatedStyle(styleFactory);

Only the function baz will be autoworkletized. Keep in mind that this is just an edge-case scenario. Please don't ever write such code when using worklets!

Scoping

We also support multiple scoping, for example:

function foo(){
  const styleFactory = () => {
    // Some UI relevant code.
  }
  function bar(){
    const animatedStyle = useAnimatedStyle(styleFactory);
  }
}

Will work as expected. It follows the same rules as above. For now we don't support variable shadowing - expect undefined behavior there.

Notes

Currently we expect the worklet variable not to be reassigned after it's been used. For example, the following code will not work:

const styleFactory = () => {
  // Some UI relevant code.
}

const animatedStyle = useAnimatedStyle(styleFactory);

styleFactory = () => {
  // Some UI relevant code.
}

Because only the last assignment will be workletized. The first assignment will not be transformed and useAnimatedStyle will throw an error. This is desired behavior, since reassigning worklet variables is considered an anti-pattern right now.

Test plan

• Add tests for all the above cases.
• Confirm that current errors remain informative in cases where this mechanism doesn't apply.

[bartlomiejbloniarz]

🚀