diff --git a/__tests__/__snapshots__/index.test.js.snap b/__tests__/__snapshots__/index.test.js.snap
new file mode 100644
index 0000000..32bcf53
--- /dev/null
+++ b/__tests__/__snapshots__/index.test.js.snap
@@ -0,0 +1,464 @@
+// Jest Snapshot v1, https://jestjs.io/docs/snapshot-testing
+
+exports[`babel-plugin-worklet basic worklet function transformation should transform arrow function with worklet directive 1`] = `
+"const myWorklet = function () {
+  const _f = function () {
+    return 42;
+  };
+  _f._closure = {};
+  _f.asString = "function _f(){return 42;}";
+  _f.__workletHash = 16000069324620;
+  _f.__location = "../test.js (1:18)";
+  _f.__worklet = true;
+  return _f;
+}();"
+`;
+
+exports[`babel-plugin-worklet basic worklet function transformation should transform function declaration with worklet directive 1`] = `
+"function myWorklet() {
+  return 42;
+}
+Object.assign(myWorklet, {
+  _closure: {},
+  asString: "function myWorklet(){return 42;}",
+  __workletHash: 9856503641981,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet basic worklet function transformation should transform function expression with worklet directive 1`] = `
+"const myWorklet = function () {
+  const _f = function () {
+    return 42;
+  };
+  _f._closure = {};
+  _f.asString = "function _f(){return 42;}";
+  _f.__workletHash = 16000069324620;
+  _f.__location = "../test.js (1:18)";
+  _f.__worklet = true;
+  return _f;
+}();"
+`;
+
+exports[`babel-plugin-worklet basic worklet function transformation should transform named function expression with worklet directive 1`] = `
+"const myWorklet = function () {
+  const _f = function () {
+    return 42;
+  };
+  _f._closure = {};
+  _f.asString = "function namedFunc(){return 42;}";
+  _f.__workletHash = 8194838852680;
+  _f.__location = "../test.js (1:18)";
+  _f.__worklet = true;
+  return _f;
+}();"
+`;
+
+exports[`babel-plugin-worklet closure variable capture should capture multiple external variables 1`] = `
+"const a = 1;
+const b = 2;
+function myWorklet() {
+  return a + b;
+}
+Object.assign(myWorklet, {
+  _closure: {
+    a,
+    b
+  },
+  asString: "function myWorklet(){const{a,b}=jsThis._closure;{return a+b;}}",
+  __workletHash: 7586383649504,
+  __location: "../test.js (3:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet closure variable capture should capture object property access 1`] = `
+"const obj = {
+  value: 10
+};
+function myWorklet() {
+  return obj.value;
+}
+Object.assign(myWorklet, {
+  _closure: {
+    obj
+  },
+  asString: "function myWorklet(){const{obj}=jsThis._closure;{return obj.value;}}",
+  __workletHash: 11513419606786,
+  __location: "../test.js (2:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet closure variable capture should capture single external variable 1`] = `
+"const x = 10;
+function myWorklet() {
+  return x + 1;
+}
+Object.assign(myWorklet, {
+  _closure: {
+    x
+  },
+  asString: "function myWorklet(){const{x}=jsThis._closure;{return x+1;}}",
+  __workletHash: 15645308083485,
+  __location: "../test.js (2:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet closure variable capture should not capture function parameters 1`] = `
+"function myWorklet(x, y) {
+  return x + y;
+}
+Object.assign(myWorklet, {
+  _closure: {},
+  asString: "function myWorklet(x,y){return x+y;}",
+  __workletHash: 12411739907900,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet closure variable capture should not capture global variables 1`] = `
+"function myWorklet() {
+  console.log(Math.PI);
+  return Date.now();
+}
+Object.assign(myWorklet, {
+  _closure: {},
+  asString: "function myWorklet(){console.log(Math.PI);return Date.now();}",
+  __workletHash: 242891576374,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet complex scenarios should handle .value assignment 1`] = `
+"const sharedValue = {
+  value: 0
+};
+function myWorklet() {
+  sharedValue.value = 100;
+}
+Object.assign(myWorklet, {
+  _closure: {
+    sharedValue
+  },
+  asString: "function myWorklet(){const{sharedValue}=jsThis._closure;{sharedValue.value=100;}}",
+  __workletHash: 15318093166788,
+  __location: "../test.js (2:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet complex scenarios should handle function with default parameters 1`] = `
+"function myWorklet(x = 10) {
+  return x * 2;
+}
+Object.assign(myWorklet, {
+  _closure: {},
+  asString: "function myWorklet(x=10){return x*2;}",
+  __workletHash: 8506619698271,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet complex scenarios should handle function with rest parameters 1`] = `
+"function myWorklet(...args) {
+  return args.length;
+}
+Object.assign(myWorklet, {
+  _closure: {},
+  asString: "function myWorklet(...args){return args.length;}",
+  __workletHash: 130556365703,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet complex scenarios should handle multiple worklet functions 1`] = `
+"function worklet1() {
+  return 1;
+}
+Object.assign(worklet1, {
+  _closure: {},
+  asString: "function worklet1(){return 1;}",
+  __workletHash: 6245207069967,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});
+function worklet2() {
+  return 2;
+}
+Object.assign(worklet2, {
+  _closure: {},
+  asString: "function worklet2(){return 2;}",
+  __workletHash: 3020052792815,
+  __location: "../test.js (5:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet custom globals should capture external variable when not configured as global 1`] = `
+"function myWorklet() {
+  return customGlobal + 1;
+}
+Object.assign(myWorklet, {
+  _closure: {
+    customGlobal
+  },
+  asString: "function myWorklet(){const{customGlobal}=jsThis._closure;{return customGlobal+1;}}",
+  __workletHash: 14087757132861,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet custom globals should not capture variable when configured as global 1`] = `
+"function myWorklet() {
+  return customGlobal + 1;
+}
+Object.assign(myWorklet, {
+  _closure: {},
+  asString: "function myWorklet(){return customGlobal+1;}",
+  __workletHash: 14354558759161,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet nested functions should handle worklet with inner function 1`] = `
+"function outer() {
+  function inner() {
+    return 1;
+  }
+  return inner();
+}
+Object.assign(outer, {
+  _closure: {},
+  asString: "function outer(){function inner(){return 1;}return inner();}",
+  __workletHash: 6466997960846,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet object property worklet should generate factory for arrow function in object property 1`] = `
+"const obj = {
+  myMethod: () => {
+    return 42;
+  },
+  _myMethod_worklet_factory_: function () {
+    var f = function () {
+      const _f = function () {
+        return 42;
+      };
+      _f._closure = {};
+      _f.asString = "function _f(){return 42;}";
+      _f.__workletHash = 16000069324620;
+      _f.__location = "../test.js";
+      _f.__worklet = true;
+      return _f;
+    }();
+    return f;
+  }
+};"
+`;
+
+exports[`babel-plugin-worklet object property worklet should generate factory for function expression in object property 1`] = `
+"const obj = {
+  myMethod: function () {
+    return 42;
+  },
+  _myMethod_worklet_factory_: function () {
+    var f = function () {
+      const _f = function () {
+        return 42;
+      };
+      _f._closure = {};
+      _f.asString = "function _f(){return 42;}";
+      _f.__workletHash = 16000069324620;
+      _f.__location = "../test.js";
+      _f.__worklet = true;
+      return _f;
+    }();
+    return f;
+  }
+};"
+`;
+
+exports[`babel-plugin-worklet object property worklet should generate factory for object method 1`] = `
+"const obj = {
+  myMethod() {
+    'worklet';
+
+    return 42;
+  }
+};"
+`;
+
+exports[`babel-plugin-worklet real world usage should handle component options with worklet methods 1`] = `
+"const option = {
+  onload() {
+    this.offset = shared(1);
+  },
+  methods: {
+    handleGesture(evt) {
+      'worklet';
+
+      runOnJs(this.bar.bind(this))();
+      this.foo();
+    },
+    bar() {
+      this.offset = 2;
+    },
+    foo(evt) {
+      'worklet';
+
+      this.offset.value = evt.translateY;
+    }
+  }
+};"
+`;
+
+exports[`babel-plugin-worklet real world usage should handle shared, derived and runOnUI pattern 1`] = `
+"const offset = shared(1);
+const func = function () {
+  const _f = function () {
+    offset.value = 2;
+  };
+  _f._closure = {
+    offset
+  };
+  _f.asString = "function _f(){const{offset}=jsThis._closure;{offset.value=2;}}";
+  _f.__workletHash = 12972093810006;
+  _f.__location = "../test.js (2:13)";
+  _f.__worklet = true;
+  return _f;
+}();
+const offset2 = derived(function () {
+  const _f = function () {
+    return offset.value * 2;
+  };
+  _f._closure = {
+    offset
+  };
+  _f.asString = "function _f(){const{offset}=jsThis._closure;{return offset.value*2;}}";
+  _f.__workletHash = 1164286910923;
+  _f.__location = "../test.js (6:24)";
+  _f.__worklet = true;
+  return _f;
+}());
+runOnUI(function () {
+  const _f = function () {
+    func();
+  };
+  _f._closure = {
+    func
+  };
+  _f.asString = "function _f(){const{func}=jsThis._closure;{func();}}";
+  _f.__workletHash = 14059708939421;
+  _f.__location = "../test.js (10:8)";
+  _f.__worklet = true;
+  return _f;
+}());"
+`;
+
+exports[`babel-plugin-worklet should not transform functions without worklet directive normal arrow function should remain unchanged 1`] = `
+"const normalFunc = () => {
+  return 42;
+};"
+`;
+
+exports[`babel-plugin-worklet should not transform functions without worklet directive normal function declaration should remain unchanged 1`] = `
+"function normalFunc() {
+  return 42;
+}"
+`;
+
+exports[`babel-plugin-worklet should not transform functions without worklet directive normal function expression should remain unchanged 1`] = `
+"const normalFunc = function () {
+  return 42;
+};"
+`;
+
+exports[`babel-plugin-worklet special syntax support should support array destructuring parameters 1`] = `
+"function myWorklet([a, b]) {
+  return a + b;
+}
+Object.assign(myWorklet, {
+  _closure: {},
+  asString: "function myWorklet([a,b]){return a+b;}",
+  __workletHash: 8542270735258,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet special syntax support should support destructuring parameters 1`] = `
+"function myWorklet({
+  a,
+  b
+}) {
+  return a + b;
+}
+Object.assign(myWorklet, {
+  _closure: {},
+  asString: "function myWorklet({a:a,b:b}){return a+b;}",
+  __workletHash: 7278292966937,
+  __location: "../test.js (1:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet special syntax support should support nullish coalescing 1`] = `
+"const val = null;
+function myWorklet() {
+  return val ?? 'default';
+}
+Object.assign(myWorklet, {
+  _closure: {
+    val
+  },
+  asString: "function myWorklet(){const{val}=jsThis._closure;{var _val;return(_val=val)!==null&&_val!==void 0?_val:'default';}}",
+  __workletHash: 1668571165844,
+  __location: "../test.js (2:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet special syntax support should support optional chaining 1`] = `
+"const obj = {
+  a: 1
+};
+function myWorklet() {
+  return obj?.a;
+}
+Object.assign(myWorklet, {
+  _closure: {
+    obj
+  },
+  asString: "function myWorklet(){const{obj}=jsThis._closure;{var _obj;return(_obj=obj)===null||_obj===void 0?void 0:_obj.a;}}",
+  __workletHash: 6861070889044,
+  __location: "../test.js (2:0)",
+  __worklet: true
+});"
+`;
+
+exports[`babel-plugin-worklet special syntax support should support template literals 1`] = `
+"const name = 'world';
+function myWorklet() {
+  return \`hello \${name}\`;
+}
+Object.assign(myWorklet, {
+  _closure: {
+    name
+  },
+  asString: "function myWorklet(){const{name}=jsThis._closure;{return\\"hello \\"+name;}}",
+  __workletHash: 14215567640686,
+  __location: "../test.js (2:0)",
+  __worklet: true
+});"
+`;
diff --git a/__tests__/index.test.js b/__tests__/index.test.js
new file mode 100644
index 0000000..d5465b8
--- /dev/null
+++ b/__tests__/index.test.js
@@ -0,0 +1,322 @@
+'use strict'
+
+const babel = require('@babel/core')
+const plugin = require('../lib/index')
+
+function transform(code, filename = 'test.js') {
+  const result = babel.transformSync(code, {
+    filename,
+    plugins: [plugin],
+    babelrc: false,
+    configFile: false,
+  })
+  return result.code
+}
+
+function transformWithOptions(code, options, filename = 'test.js') {
+  const result = babel.transformSync(code, {
+    filename,
+    plugins: [[plugin, options]],
+    babelrc: false,
+    configFile: false,
+  })
+  return result.code
+}
+
+describe('babel-plugin-worklet', () => {
+  describe('basic worklet function transformation', () => {
+    test('should transform function declaration with worklet directive', () => {
+      const input = `function myWorklet() {
+  'worklet';
+  return 42;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should transform arrow function with worklet directive', () => {
+      const input = `const myWorklet = () => {
+  'worklet';
+  return 42;
+};`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should transform function expression with worklet directive', () => {
+      const input = `const myWorklet = function() {
+  'worklet';
+  return 42;
+};`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should transform named function expression with worklet directive', () => {
+      const input = `const myWorklet = function namedFunc() {
+  'worklet';
+  return 42;
+};`
+      expect(transform(input)).toMatchSnapshot()
+    })
+  })
+
+  describe('closure variable capture', () => {
+    test('should capture single external variable', () => {
+      const input = `const x = 10;
+function myWorklet() {
+  'worklet';
+  return x + 1;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should capture multiple external variables', () => {
+      const input = `const a = 1;
+const b = 2;
+function myWorklet() {
+  'worklet';
+  return a + b;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should not capture global variables', () => {
+      const input = `function myWorklet() {
+  'worklet';
+  console.log(Math.PI);
+  return Date.now();
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should not capture function parameters', () => {
+      const input = `function myWorklet(x, y) {
+  'worklet';
+  return x + y;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should capture object property access', () => {
+      const input = `const obj = { value: 10 };
+function myWorklet() {
+  'worklet';
+  return obj.value;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+  })
+
+  describe('object property worklet', () => {
+    test('should generate factory for function expression in object property', () => {
+      const input = `const obj = {
+  myMethod: function() {
+    'worklet';
+    return 42;
+  }
+};`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should generate factory for arrow function in object property', () => {
+      const input = `const obj = {
+  myMethod: () => {
+    'worklet';
+    return 42;
+  }
+};`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should generate factory for object method', () => {
+      const input = `const obj = {
+  myMethod() {
+    'worklet';
+    return 42;
+  }
+};`
+      expect(transform(input)).toMatchSnapshot()
+    })
+  })
+
+  describe('should not transform functions without worklet directive', () => {
+    test('normal function declaration should remain unchanged', () => {
+      const input = `function normalFunc() {
+  return 42;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('normal arrow function should remain unchanged', () => {
+      const input = `const normalFunc = () => {
+  return 42;
+};`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('normal function expression should remain unchanged', () => {
+      const input = `const normalFunc = function() {
+  return 42;
+};`
+      expect(transform(input)).toMatchSnapshot()
+    })
+  })
+
+  describe('custom globals', () => {
+    test('should capture external variable when not configured as global', () => {
+      const input = `function myWorklet() {
+  'worklet';
+  return customGlobal + 1;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should not capture variable when configured as global', () => {
+      const input = `function myWorklet() {
+  'worklet';
+  return customGlobal + 1;
+}`
+      expect(transformWithOptions(input, { globals: ['customGlobal'] })).toMatchSnapshot()
+    })
+  })
+
+  describe('nested functions', () => {
+    test('should handle worklet with inner function', () => {
+      const input = `function outer() {
+  'worklet';
+  function inner() {
+    return 1;
+  }
+  return inner();
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+  })
+
+  describe('complex scenarios', () => {
+    test('should handle multiple worklet functions', () => {
+      const input = `function worklet1() {
+  'worklet';
+  return 1;
+}
+function worklet2() {
+  'worklet';
+  return 2;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should handle .value assignment', () => {
+      const input = `const sharedValue = { value: 0 };
+function myWorklet() {
+  'worklet';
+  sharedValue.value = 100;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should handle function with default parameters', () => {
+      const input = `function myWorklet(x = 10) {
+  'worklet';
+  return x * 2;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should handle function with rest parameters', () => {
+      const input = `function myWorklet(...args) {
+  'worklet';
+  return args.length;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+  })
+
+  describe('special syntax support', () => {
+    test('should support optional chaining', () => {
+      const input = `const obj = { a: 1 };
+function myWorklet() {
+  'worklet';
+  return obj?.a;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should support nullish coalescing', () => {
+      const input = `const val = null;
+function myWorklet() {
+  'worklet';
+  return val ?? 'default';
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should support template literals', () => {
+      const input = `const name = 'world';
+function myWorklet() {
+  'worklet';
+  return \`hello \${name}\`;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should support destructuring parameters', () => {
+      const input = `function myWorklet({ a, b }) {
+  'worklet';
+  return a + b;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should support array destructuring parameters', () => {
+      const input = `function myWorklet([a, b]) {
+  'worklet';
+  return a + b;
+}`
+      expect(transform(input)).toMatchSnapshot()
+    })
+  })
+
+  describe('real world usage', () => {
+    test('should handle shared, derived and runOnUI pattern', () => {
+      const input = `const offset = shared(1)
+const func = () => {
+  'worklet';
+  offset.value = 2
+}
+const offset2 = derived(() => {
+  'worklet';
+  return offset.value * 2
+})
+runOnUI(() => {
+  'worklet';
+  func()
+})`
+      expect(transform(input)).toMatchSnapshot()
+    })
+
+    test('should handle component options with worklet methods', () => {
+      const input = `const option = {
+  onload() {
+    this.offset = shared(1)
+  },
+ 
+  methods: {
+    handleGesture(evt) {
+      'worklet';
+      runOnJs(this.bar.bind(this))()
+      this.foo()
+    },
+
+    bar() {
+      this.offset = 2
+    },
+
+    foo(evt) {
+      'worklet';
+      this.offset.value = evt.translateY
+    }
+  }
+};`
+      expect(transform(input)).toMatchSnapshot()
+    })
+  })
+})
diff --git a/lib/index.js b/lib/index.js
index 50a4db7..aab6d8e 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -351,16 +351,13 @@ function makeWorkletName(t, fun) {
   return '_f' // fallback for ArrowFunctionExpression and unnamed FunctionExpression
 }
 
-function makeWorklet(t, fun, fileName) {
-  // Returns a new FunctionExpression which is a workletized version of provided
-  // FunctionDeclaration, FunctionExpression, ArrowFunctionExpression or ObjectMethod.
+function extractWorkletData(t, fun, fileName) {
+  // Extract worklet data: closure, variables, funString, workletHash, locationFileName
+  // This is the common logic shared by makeWorklet and buildWorkletProperties
 
   const functionName = makeWorkletName(t, fun)
-
   const closure = new Map()
-  const outputs = new Set()
   const closureGenerator = new ClosureGenerator()
-  const options = {}
 
   // remove 'worklet'; directive before calling .toString()
   fun.traverse({
@@ -373,7 +370,6 @@ function makeWorklet(t, fun, fileName) {
 
   // We use copy because some of the plugins don't update bindings and
   // some even break them
-
   const code = '\n(' + (t.isObjectMethod(fun) ? 'function ' : '') + fun.toString() + '\n)'
 
   const transformed = transformSync(code, {
@@ -390,9 +386,7 @@ function makeWorklet(t, fun, fileName) {
     babelrc: false,
     configFile: false,
   })
-  if (fun.parent && fun.parent.callee && fun.parent.callee.name === 'createAnimatedStyle') {
-    options.optFlags = isPossibleOptimization(transformed.ast)
-  }
+
   traverse(transformed.ast, {
     ReferencedIdentifier(path) {
       const name = path.node.name
@@ -425,40 +419,50 @@ function makeWorklet(t, fun, fileName) {
       closure.set(name, path.node)
       closureGenerator.addPath(name, path)
     },
-    AssignmentExpression(path) {
-      // test for <something>.value = <something> expressions
-      const left = path.node.left
-      if (
-        t.isMemberExpression(left) &&
-        t.isIdentifier(left.object) &&
-        t.isIdentifier(left.property, { name: 'value' })
-      ) {
-        outputs.add(left.object.name)
-      }
-    },
   })
 
   const variables = Array.from(closure.values())
-
-  const privateFunctionId = t.identifier('_f')
-  const clone = t.cloneNode(fun.node)
-  let funExpression
-  if (clone.body.type === 'BlockStatement') {
-    funExpression = t.functionExpression(null, clone.params, clone.body)
-  } else {
-    funExpression = clone
-  }
   const funString = buildWorkletString(t, transformed.ast, variables, functionName)
   const workletHash = hash(funString)
 
+  let locationFileName = fileName
   const loc = fun && fun.node && fun.node.loc && fun.node.loc.start
   if (loc) {
     const { line, column } = loc
     if (typeof line === 'number' && typeof column === 'number') {
-      fileName = `${fileName} (${line}:${column})`
+      locationFileName = `${fileName} (${line}:${column})`
     }
   }
 
+  return {
+    variables,
+    closure,
+    closureGenerator,
+    funString,
+    workletHash,
+    locationFileName,
+  }
+}
+
+function makeWorklet(t, fun, fileName) {
+  // Returns a new FunctionExpression which is a workletized version of provided
+  // FunctionDeclaration, FunctionExpression, ArrowFunctionExpression or ObjectMethod.
+
+  const { variables, closure, closureGenerator, funString, workletHash, locationFileName } = extractWorkletData(
+    t,
+    fun,
+    fileName,
+  )
+
+  const privateFunctionId = t.identifier('_f')
+  const clone = t.cloneNode(fun.node)
+  let funExpression
+  if (clone.body.type === 'BlockStatement') {
+    funExpression = t.functionExpression(null, clone.params, clone.body)
+  } else {
+    funExpression = clone
+  }
+
   const statements = [
     t.variableDeclaration('const', [t.variableDeclarator(privateFunctionId, funExpression)]),
     t.expressionStatement(
@@ -486,7 +490,7 @@ function makeWorklet(t, fun, fileName) {
       t.assignmentExpression(
         '=',
         t.memberExpression(privateFunctionId, t.identifier('__location'), false),
-        t.stringLiteral(fileName),
+        t.stringLiteral(locationFileName),
       ),
     ),
     t.expressionStatement(
@@ -505,6 +509,24 @@ function makeWorklet(t, fun, fileName) {
   return newFun
 }
 
+function buildWorkletProperties(t, fun, fileName) {
+  // Build the worklet properties object for Object.assign
+
+  const { variables, closure, closureGenerator, funString, workletHash, locationFileName } = extractWorkletData(
+    t,
+    fun,
+    fileName,
+  )
+
+  return t.objectExpression([
+    t.objectProperty(t.identifier('_closure'), closureGenerator.generate(t, variables, closure.keys())),
+    t.objectProperty(t.identifier('asString'), t.stringLiteral(funString)),
+    t.objectProperty(t.identifier('__workletHash'), t.numericLiteral(workletHash)),
+    t.objectProperty(t.identifier('__location'), t.stringLiteral(locationFileName)),
+    t.objectProperty(t.identifier('__worklet'), t.booleanLiteral(true)),
+  ])
+}
+
 function processWorkletFunction(t, fun, fileName) {
   // Replaces FunctionDeclaration, FunctionExpression or ArrowFunctionExpression
   // with a workletized version of itself.
@@ -539,21 +561,41 @@ function processWorkletFunction(t, fun, fileName) {
     return
   }
 
-  const newFun = makeWorklet(t, fun, fileName)
-
-  const replacement = t.callExpression(newFun, [])
-
-  // we check if function needs to be assigned to variable declaration.
-  // This is needed if function definition directly in a scope. Some other ways
-  // where function definition can be used is for example with variable declaration:
-  // const ggg = function foo() { }
-  // ^ in such a case we don't need to define variable for the function
+  // Check if this is a FunctionDeclaration at scope level (needs hoisting preservation)
   const needDeclaration = t.isScopable(fun.parent) || t.isExportNamedDeclaration(fun.parent)
-  fun.replaceWith(
-    fun.node.id && needDeclaration
-      ? t.variableDeclaration('const', [t.variableDeclarator(fun.node.id, replacement)])
-      : replacement,
-  )
+  const isFunctionDeclaration = t.isFunctionDeclaration(fun.node) && fun.node.id && needDeclaration
+
+  if (isFunctionDeclaration) {
+    // For function declarations, keep the function and use Object.assign to add properties
+    const functionId = fun.node.id
+    const properties = buildWorkletProperties(t, fun, fileName)
+
+    // Create Object.assign statement
+    const objectAssignCall = t.expressionStatement(
+      t.callExpression(t.memberExpression(t.identifier('Object'), t.identifier('assign')), [
+        t.identifier(functionId.name),
+        properties,
+      ]),
+    )
+
+    // Insert Object.assign after the function declaration
+    fun.insertAfter(objectAssignCall)
+  } else {
+    // For other cases (arrow functions, function expressions), use the original IIFE approach
+    const newFun = makeWorklet(t, fun, fileName)
+    const replacement = t.callExpression(newFun, [])
+
+    // we check if function needs to be assigned to variable declaration.
+    // This is needed if function definition directly in a scope. Some other ways
+    // where function definition can be used is for example with variable declaration:
+    // const ggg = function foo() { }
+    // ^ in such a case we don't need to define variable for the function
+    fun.replaceWith(
+      fun.node.id && needDeclaration
+        ? t.variableDeclaration('const', [t.variableDeclarator(fun.node.id, replacement)])
+        : replacement,
+    )
+  }
 }
 
 function processIfWorkletNode(t, fun, fileName) {
diff --git a/package.json b/package.json
index bcec3b3..bec6107 100644
--- a/package.json
+++ b/package.json
@@ -4,7 +4,7 @@
   "description": "Compile worklet function in miniprogram",
   "main": "lib/index",
   "scripts": {
-    "test": "echo \"Error: no test specified\" && exit 1"
+    "test": "jest"
   },
   "repository": {
     "type": "git",
@@ -21,9 +21,15 @@
   },
   "homepage": "https://github.com/wechat-miniprogram/babel-plugin-worklet#readme",
   "devDependencies": {
+    "@babel/core": "^7.28.5",
     "@babel/plugin-proposal-class-properties": "^7.18.6",
+    "@babel/plugin-proposal-nullish-coalescing-operator": "^7.18.6",
+    "@babel/plugin-proposal-optional-chaining": "^7.21.0",
     "@babel/plugin-transform-arrow-functions": "^7.18.6",
     "@babel/plugin-transform-shorthand-properties": "^7.18.6",
+    "@babel/plugin-transform-template-literals": "^7.27.1",
+    "@babel/preset-typescript": "^7.28.5",
+    "jest": "^30.2.0",
     "string-hash-64": "^1.0.3"
   }
-}
\ No newline at end of file
+}