diff --git a/packages/define-render/package.json b/packages/define-render/package.json
index ecd5e4403..af429c1aa 100644
--- a/packages/define-render/package.json
+++ b/packages/define-render/package.json
@@ -89,6 +89,7 @@
     "vue": "^2.7.0 || ^3.0.0"
   },
   "dependencies": {
+    "@ast-grep/napi": "^0.5.7",
     "@vue-macros/common": "workspace:~",
     "unplugin": "^1.3.1"
   },
diff --git a/packages/define-render/src/core/index.ts b/packages/define-render/src/core/index.ts
index e25c9881b..fa7e2aa5d 100644
--- a/packages/define-render/src/core/index.ts
+++ b/packages/define-render/src/core/index.ts
@@ -1,67 +1,66 @@
 import {
   DEFINE_RENDER,
   MagicString,
-  babelParse,
   getLang,
   getTransformResult,
-  isCallOf,
-  isFunctionType,
-  walkAST,
 } from '@vue-macros/common'
-import {
-  type BlockStatement,
-  type ExpressionStatement,
-  type Node,
-} from '@babel/types'
+import { js, jsx, ts, tsx } from '@ast-grep/napi'
+
+const isFunction = (kind: string) => kind.includes('function')
 
-export function transformDefineRender(code: string, id: string) {
+export const transformDefineRender = (code: string, id: string) => {
   if (!code.includes(DEFINE_RENDER)) return
 
   const lang = getLang(id)
-  const program = babelParse(code, lang === 'vue' ? 'js' : lang)
 
-  const nodes: {
-    parent: BlockStatement
-    node: ExpressionStatement
-    arg: Node
-  }[] = []
-  walkAST<Node>(program, {
-    enter(node, parent) {
-      if (
-        node.type !== 'ExpressionStatement' ||
-        !isCallOf(node.expression, DEFINE_RENDER) ||
-        parent?.type !== 'BlockStatement'
-      )
-        return
+  const processor = { js, jsx, ts, tsx, vue: jsx }[lang]
+  if (!processor) {
+    throw new Error(`not supported lang: ${lang}`)
+  }
+  const root = processor.parse(code).root()
 
-      nodes.push({
-        parent,
-        node,
-        arg: node.expression.arguments[0],
-      })
-    },
-  })
+  const nodes = root.findAll('defineRender($$$)')
   if (nodes.length === 0) return
 
   const s = new MagicString(code)
 
-  for (const { parent, node, arg } of nodes) {
-    // check parent
-    const returnStmt = parent.body.find(
-      (node) => node.type === 'ReturnStatement'
-    )
-    if (returnStmt) s.removeNode(returnStmt)
+  for (const node of nodes) {
+    const args = node.field('arguments')?.children()
+    if (!args || args.length < 3 /* '(', first arg, ')' */)
+      throw new SyntaxError(`bad arguments: ${node.text()}`)
+
+    const [, arg] = args
+    const argRng = arg.range()
+
+    const parents = node.ancestors()
+    if (parents[0].kind() !== 'expression_statement' || !parents[1]) {
+      throw new SyntaxError(`bad case: ${node.text()}`)
+    }
+
+    // remove ReturnStatement of the parent
+    const returnStmtRange = parents[1]
+      .children()
+      .find((n) => n.kind() === 'return_statement')
+      ?.range()
+
+    if (returnStmtRange)
+      s.remove(returnStmtRange.start.index, returnStmtRange.end.index)
+
+    const lastChild = parents[1].children().at(-1)!
+    const index = returnStmtRange
+      ? returnStmtRange.start.index
+      : lastChild.range()[lastChild.kind() === '}' ? 'start' : 'end'].index
 
-    const index = returnStmt ? returnStmt.start! : parent.end! - 1
-    const shouldAddFn = !isFunctionType(arg) && arg.type !== 'Identifier'
+    const shouldAddFn = !isFunction(arg.kind()) && arg.kind() !== 'identifier'
     s.appendLeft(index, `return ${shouldAddFn ? '() => (' : ''}`)
-    s.moveNode(arg, index)
+    s.move(argRng.start.index, argRng.end.index, index)
     if (shouldAddFn) s.appendRight(index, `)`)
 
+    const nodeRng = node.range()
     // removes `defineRender(`
-    s.remove(node.start!, arg.start!)
+    s.remove(nodeRng.start.index, argRng.start.index)
     // removes `)`
-    s.remove(arg.end!, node.end!)
+    s.remove(argRng.end.index, parents[0].range().end.index)
   }
 
   return getTransformResult(s, id)
diff --git a/packages/define-render/tests/__snapshots__/fixtures.test.ts.snap b/packages/define-render/tests/__snapshots__/fixtures.test.ts.snap
index abdb7743e..8a6e473db 100644
--- a/packages/define-render/tests/__snapshots__/fixtures.test.ts.snap
+++ b/packages/define-render/tests/__snapshots__/fixtures.test.ts.snap
@@ -1,5 +1,16 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
+exports[`fixtures > tests/fixtures/basic.jsx 1`] = `
+"import { defineComponent, createVNode } from 'vue';
+
+/* @__PURE__ */ defineComponent({
+  setup() {
+    defineRender(createVNode(\\"div\\", null, null));
+  }
+});
+"
+`;
+
 exports[`fixtures > tests/fixtures/basic.vue 1`] = `
 "import { defineComponent, h } from 'vue';
 import _export_sfc from '[NULL]/plugin-vue/export-helper';
@@ -17,6 +28,8 @@ export { basic as default };
 "
 `;
 
+exports[`fixtures > tests/fixtures/error1.vue 1`] = `"bad arguments: defineRender()"`;
+
 exports[`fixtures > tests/fixtures/id.vue 1`] = `
 "import { defineComponent, createVNode, createTextVNode } from 'vue';
 import _export_sfc from '[NULL]/plugin-vue/export-helper';
@@ -52,6 +65,18 @@ export { tsx as default };
 "
 `;
 
+exports[`fixtures > tests/fixtures/with-return.js 1`] = `
+";
+() => {
+  return () => {
+  };
+  defineRender(() => {
+    console.log(\\"hello world\\");
+  });
+};
+"
+`;
+
 exports[`fixtures > tests/fixtures/without-fn.vue 1`] = `
 "import { defineComponent, createVNode, createTextVNode } from 'vue';
 import _export_sfc from '[NULL]/plugin-vue/export-helper';
diff --git a/packages/define-render/tests/fixtures.test.ts b/packages/define-render/tests/fixtures.test.ts
index 5c8c893f3..0aaa95d1a 100644
--- a/packages/define-render/tests/fixtures.test.ts
+++ b/packages/define-render/tests/fixtures.test.ts
@@ -11,7 +11,7 @@ import VueDefineRender from '../src/rollup'
 
 describe('fixtures', async () => {
   await testFixtures(
-    'tests/fixtures/*.{vue,js,ts}',
+    'tests/fixtures/*.{vue,[jt]s?(x)}',
     (args, id) =>
       rollupBuild(id, [
         RollupVue(),
diff --git a/packages/define-render/tests/fixtures/basic.jsx b/packages/define-render/tests/fixtures/basic.jsx
new file mode 100644
index 000000000..21805fbfc
--- /dev/null
+++ b/packages/define-render/tests/fixtures/basic.jsx
@@ -0,0 +1,7 @@
+import { defineComponent } from 'vue'
+
+defineComponent({
+  setup() {
+    defineRender(<div />)
+  },
+})
diff --git a/packages/define-render/tests/fixtures/error1.vue b/packages/define-render/tests/fixtures/error1.vue
new file mode 100644
index 000000000..b6422203a
--- /dev/null
+++ b/packages/define-render/tests/fixtures/error1.vue
@@ -0,0 +1,3 @@
+<script setup lang="tsx">
+defineRender()
+</script>
diff --git a/packages/define-render/tests/fixtures/with-return.js b/packages/define-render/tests/fixtures/with-return.js
new file mode 100644
index 000000000..04a145c0e
--- /dev/null
+++ b/packages/define-render/tests/fixtures/with-return.js
@@ -0,0 +1,7 @@
+;() => {
+  return () => {}
+
+  defineRender(() => {
+    console.log('hello world')
+  })
+}
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 4029520be..a7fa2c94b 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -284,6 +284,9 @@ importers:
 
   packages/define-render:
     dependencies:
+      '@ast-grep/napi':
+        specifier: ^0.5.7
+        version: 0.5.7
       '@vue-macros/common':
         specifier: workspace:~
         version: link:../common
@@ -907,6 +910,72 @@ packages:
       leven: 3.1.0
     dev: true
 
+  /@ast-grep/napi-darwin-arm64@0.5.7:
+    resolution: {integrity: sha512-f9oQBA527T6Wx7P/apACsbDym9rZyy0zMCzrYgyxGEw9owoDL6Z+nWc5JDd7kM98zhaUBzoBHrls8xtv/xT49Q==}
+    engines: {node: '>= 10'}
+    cpu: [arm64]
+    os: [darwin]
+    requiresBuild: true
+    dev: false
+    optional: true
+
+  /@ast-grep/napi-darwin-x64@0.5.7:
+    resolution: {integrity: sha512-bphqeNXpxDB1CWpBEjmvFYLMRmmipmCBuQ76TjOglCdT+Pah/kJ76mCyRaYHLhW104M/9XtBi82OefS1U8OrkA==}
+    engines: {node: '>= 10'}
+    cpu: [x64]
+    os: [darwin]
+    requiresBuild: true
+    dev: false
+    optional: true
+
+  /@ast-grep/napi-linux-x64-gnu@0.5.7:
+    resolution: {integrity: sha512-vpvVQ3F+2OZH3J5Au1yIijFmFVmRrYRrW2F0BW2VaBoHc+5xyCmCeEZVfitn4tZcN2WBTLaFeh2IAFDOWp1XjA==}
+    engines: {node: '>= 10'}
+    cpu: [x64]
+    os: [linux]
+    requiresBuild: true
+    dev: false
+    optional: true
+
+  /@ast-grep/napi-win32-arm64-msvc@0.5.7:
+    resolution: {integrity: sha512-2e8QkdMISdC/JQJ/Hm1KjWPls4saiiMF4mfpK6gUNAe/Kqy9yBAwv8YdCzz6ucO0JHR2056zAGaJvG8YgcHDtg==}
+    engines: {node: '>= 10'}
+    cpu: [arm64]
+    os: [win32]
+    requiresBuild: true
+    dev: false
+    optional: true
+
+  /@ast-grep/napi-win32-ia32-msvc@0.5.7:
+    resolution: {integrity: sha512-7wOd7O3ozWkXhUVlOvxM13IPl98TefhacH9ZbriPtOFbuO+2dP+xIqh5F1CAcFfltN8r/3pZecp33hvSPccDlQ==}
+    engines: {node: '>= 10'}
+    cpu: [ia32]
+    os: [win32]
+    requiresBuild: true
+    dev: false
+    optional: true
+
+  /@ast-grep/napi-win32-x64-msvc@0.5.7:
+    resolution: {integrity: sha512-u17aQUI5TtU8xCWC23ikTSQKDPfjJyXOHNIxZ1pT0v2LCfqV0hl/iwe3LKXXGPPs2+tEZVPNawV8R1l+KLSRAg==}
+    engines: {node: '>= 10'}
+    cpu: [x64]
+    os: [win32]
+    requiresBuild: true
+    dev: false
+    optional: true
+
+  /@ast-grep/napi@0.5.7:
+    resolution: {integrity: sha512-t4XEGqQDeVIJAap6+O1SykddFtruTB79++lR2oRF83f5KesE6tcQaxeJb/jkdy5atmlUx4lOjUbtNFY489oKsQ==}
+    engines: {node: '>= 10'}
+    optionalDependencies:
+      '@ast-grep/napi-darwin-arm64': 0.5.7
+      '@ast-grep/napi-darwin-x64': 0.5.7
+      '@ast-grep/napi-linux-x64-gnu': 0.5.7
+      '@ast-grep/napi-win32-arm64-msvc': 0.5.7
+      '@ast-grep/napi-win32-ia32-msvc': 0.5.7
+      '@ast-grep/napi-win32-x64-msvc': 0.5.7
+    dev: false
+
   /@babel/code-frame@7.21.4:
     resolution: {integrity: sha512-LYvhNKfwWSPpocw8GI7gpK2nq3HSDuEPC/uSYaALSJu9xjsalaaYFOq0Pwt5KmVqwEbZlDu81aLXwBOmD/Fv9g==}
     engines: {node: '>=6.9.0'}