From eb1dd9bc799a59a9425a9c7262093c97534ada21 Mon Sep 17 00:00:00 2001
From: Justin Schroeder <justin@formkit.com>
Date: Fri, 29 Dec 2023 16:36:12 -0500
Subject: [PATCH] fix: moves provider inside root node if possible.
 (formkit/formkit#1095)

---
 src/index.ts                          | 24 +++++++++++++++---
 test/__snapshots__/index.test.ts.snap | 19 ++++++++++++---
 test/index.test.ts                    | 35 +++++++++++++++++++++------
 3 files changed, 63 insertions(+), 15 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 1ceb7ef..df7fbc5 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -10,9 +10,25 @@ function getRootBlock(
   root: RootNode,
   block: 'template' | 'script' | 'style',
 ): ElementNode | undefined {
-  return root.children.find((node) => node.type === 1 && node.tag === block) as
-    | ElementNode
-    | undefined
+  const node = root.children.find(
+    (node) => node.type === 1 && node.tag === block,
+  ) as ElementNode | undefined
+  if (node && block === 'template' && node.children.length === 1) {
+    const rootChild = node.children[0].type === 1 ? node.children[0] : undefined
+    const tag = (rootChild?.tag ?? '').toLocaleLowerCase()
+    if (
+      rootChild &&
+      tag !== 'formkit' &&
+      tag !== 'form-kit' &&
+      rootChild.isSelfClosing === false
+    ) {
+      // In this case the component has a root node that is not formkit and is
+      // not self-closing, like, perhaps, a div. We need to move the provider
+      // inside this div instead of outside it.
+      return rootChild
+    }
+  }
+  return node
 }
 
 /**
@@ -157,7 +173,7 @@ export const unpluginFactory: UnpluginFactory<Options | undefined> = (
     },
     // webpack's id filter is outside of loader logic,
     // an additional hook is needed for better perf on webpack
-    transformInclude(id: string) {
+    transformInclude(_id: string) {
       // TODO: resolve why @formkit/vue is not always identifiable by the id
       // and remove this early return workaround:
       return true
diff --git a/test/__snapshots__/index.test.ts.snap b/test/__snapshots__/index.test.ts.snap
index 15af2f3..7a1ccb3 100644
--- a/test/__snapshots__/index.test.ts.snap
+++ b/test/__snapshots__/index.test.ts.snap
@@ -17,6 +17,17 @@ import { FormKit } from '@formkit/vue'
 </template>"
 `;
 
+exports[`index > injects inside root node if there is one 1`] = `
+"<script setup>import { FormKitLazyProvider } from '@formkit/vue'</script>
+<template>
+    <div>
+    
+<FormKitLazyProvider>  <FormKit /></FormKitLazyProvider>
+
+    </div>
+  </template>"
+`;
+
 exports[`index > injects setup block when using options api 1`] = `
 "<script setup lang=\\"ts\\">import { FormKitLazyProvider } from '@formkit/vue'</script>
 <script lang=\\"ts\\">
@@ -30,12 +41,12 @@ export default {
 </script>
 
 <template>
-<FormKitLazyProvider>
   <div>
-    <h1>Nothing to see here</h1>
-    <FormKit type=\\"text\\" label=\\"Check me out\\" />
+    
+<FormKitLazyProvider><h1>Nothing to see here</h1>
+    <FormKit type=\\"text\\" label=\\"Check me out\\" </FormKitLazyProvider>
+/>
   </div>
-</FormKitLazyProvider>
 </template>"
 `;
 
diff --git a/test/index.test.ts b/test/index.test.ts
index 9050ac7..2d1c0c5 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -2,7 +2,9 @@ import { describe, expect, it } from 'vitest'
 import unplugin from '../src/vite'
 import { readFileSync } from 'fs'
 
-type Transformer = { transform: (code: string, id: string) => Promise<string> }
+type Transformer = {
+  transform: (code: string, id: string) => Promise<{ code: string; map?: any }>
+}
 const aboutSFCFile = readFileSync('./playground/src/pages/about.vue', 'utf-8')
 const contactSFCFile = readFileSync(
   './playground/src/pages/contact.vue',
@@ -14,22 +16,41 @@ const plugin: Transformer = unplugin() as Transformer
 describe('index', () => {
   it('injects the template block into an normally structured sfc', async () => {
     expect(
-      await plugin.transform(
-        `<template>
+      (
+        await plugin.transform(
+          `<template>
   <FormKit />
 </template>`,
-        'test.vue',
-      ),
+          'test.vue',
+        )
+      ).code,
+    ).toMatchSnapshot()
+  })
+
+  it('injects inside root node if there is one', async () => {
+    expect(
+      (
+        await plugin.transform(
+          `<template>
+    <div>
+      <FormKit />
+    </div>
+  </template>`,
+          'test.vue',
+        )
+      ).code,
     ).toMatchSnapshot()
   })
 
   it('injects import into script setup block', async () => {
-    expect(await plugin.transform(aboutSFCFile, 'about.vue')).toMatchSnapshot()
+    expect(
+      (await plugin.transform(aboutSFCFile, 'about.vue')).code,
+    ).toMatchSnapshot()
   })
 
   it('injects setup block when using options api', async () => {
     expect(
-      await plugin.transform(contactSFCFile, 'about.vue'),
+      (await plugin.transform(contactSFCFile, 'about.vue')).code,
     ).toMatchSnapshot()
   })
 })