@clayroach/unplugin 0.1.0-source-trace.3 → 0.1.0-source-trace.5

package/src/transformers/annotateEffects.ts

@@ -0,0 +1,197 @@
+/**
+ * Pure call annotation transformer.
+ *
+ * This transformer adds `#__PURE__` comments to call expressions for better tree-shaking.
+ * It replicates the behavior of `babel-plugin-annotate-pure-calls` but integrated into
+ * the Effect unplugin.
+ *
+ * The `#__PURE__` annotation tells bundlers (Webpack, Rollup, esbuild) that a function call
+ * has no side effects and can be safely removed if the result is unused.
+ *
+ * @since 0.1.0
+ */
+import type { NodePath, Visitor } from "@babel/traverse"
+import * as t from "@babel/types"
+
+/**
+ * Options for the annotate effects transformer.
+ */
+export interface AnnotateEffectsOptions {
+  /**
+   * Filter function to determine if a file should be transformed.
+   */
+  readonly filter?: (filename: string) => boolean
+}
+
+/**
+ * State passed through the transformer.
+ */
+interface TransformState {
+  filename: string
+}
+
+const PURE_ANNOTATION = "#__PURE__"
+
+/**
+ * Checks if a node already has a PURE annotation.
+ */
+function isPureAnnotated(node: t.Node): boolean {
+  const leadingComments = node.leadingComments
+  if (!leadingComments) {
+    return false
+  }
+  return leadingComments.some((comment) => /[@#]__PURE__/.test(comment.value))
+}
+
+/**
+ * Adds a PURE annotation comment to a node.
+ */
+function annotateAsPure(path: NodePath<t.CallExpression | t.NewExpression>): void {
+  if (isPureAnnotated(path.node)) {
+    return
+  }
+  path.addComment("leading", PURE_ANNOTATION)
+}
+
+/**
+ * Checks if the parent is a CallExpression or NewExpression.
+ */
+function hasCallableParent(path: NodePath): boolean {
+  const parentPath = path.parentPath
+  if (!parentPath) return false
+  return parentPath.isCallExpression() || parentPath.isNewExpression()
+}
+
+/**
+ * Checks if this node is used as a callee (e.g., `foo()` where foo is the callee).
+ */
+function isUsedAsCallee(path: NodePath): boolean {
+  if (!hasCallableParent(path)) {
+    return false
+  }
+  const parentPath = path.parentPath as NodePath<t.CallExpression | t.NewExpression>
+  return parentPath.get("callee") === path
+}
+
+/**
+ * Checks if this node is inside a callee chain (e.g., `foo()()` or `foo.bar()`).
+ */
+function isInCallee(path: NodePath): boolean {
+  let current: NodePath | null = path
+  do {
+    current = current.parentPath
+    if (!current) return false
+    if (isUsedAsCallee(current)) {
+      return true
+    }
+  } while (!current.isStatement() && !current.isFunction())
+  return false
+}
+
+/**
+ * Checks if this expression is executed during module initialization
+ * (not inside a function that isn't immediately invoked).
+ */
+function isExecutedDuringInitialization(path: NodePath): boolean {
+  let functionParent = path.getFunctionParent()
+  while (functionParent) {
+    if (!isUsedAsCallee(functionParent)) {
+      return false
+    }
+    functionParent = functionParent.getFunctionParent()
+  }
+  return true
+}
+
+/**
+ * Checks if this expression is in an assignment context
+ * (variable declaration, assignment expression, or class).
+ */
+function isInAssignmentContext(path: NodePath): boolean {
+  const statement = path.getStatementParent()
+  if (!statement) return false
+
+  let currentPath: NodePath | null = path
+  do {
+    currentPath = currentPath.parentPath
+    if (!currentPath) return false
+
+    if (
+      currentPath.isVariableDeclaration() ||
+      currentPath.isAssignmentExpression() ||
+      currentPath.isClass()
+    ) {
+      return true
+    }
+  } while (currentPath !== statement)
+
+  return false
+}
+
+/**
+ * Visitor function for CallExpression and NewExpression nodes.
+ */
+function callableExpressionVisitor(
+  path: NodePath<t.CallExpression | t.NewExpression>,
+  _state: TransformState
+): void {
+  // Skip if this is used as a callee (e.g., foo()())
+  if (isUsedAsCallee(path)) {
+    return
+  }
+
+  // Skip if this is inside a callee chain
+  if (isInCallee(path)) {
+    return
+  }
+
+  // Skip if not executed during initialization
+  if (!isExecutedDuringInitialization(path)) {
+    return
+  }
+
+  // Must be in assignment context or export default
+  const statement = path.getStatementParent()
+  if (!isInAssignmentContext(path) && !statement?.isExportDefaultDeclaration()) {
+    return
+  }
+
+  annotateAsPure(path)
+}
+
+/**
+ * Creates a Babel visitor that adds PURE annotations to call expressions.
+ */
+export function createAnnotateEffectsVisitor(
+  filename: string,
+  _options?: AnnotateEffectsOptions
+): Visitor<TransformState> {
+  return {
+    Program: {
+      enter(_path, state) {
+        state.filename = filename
+      }
+    },
+
+    CallExpression(path: NodePath<t.CallExpression>, state) {
+      callableExpressionVisitor(path, state)
+    },
+
+    NewExpression(path: NodePath<t.NewExpression>, state) {
+      callableExpressionVisitor(path, state)
+    }
+  }
+}
+
+/**
+ * Creates the annotate effects transformer plugin.
+ */
+export function annotateEffectsTransformer(options?: AnnotateEffectsOptions): {
+  visitor: Visitor<TransformState>
+  name: string
+} {
+  return {
+    name: "effect-annotate-pure-calls",
+    visitor: createAnnotateEffectsVisitor("", options)
+  }
+}

package/src/transformers/sourceTrace.ts

@@ -2,14 +2,23 @@
  * Source location trace transformer.
  *
  * This transformer injects source location metadata into Effect.gen yield expressions.
- * It transforms `yield* _(effect)` into `yield* _(effect, trace)` where trace is a
- * hoisted SourceLocation object.
+ * It supports both patterns:
+ * - Legacy adapter: `yield* _(effect)` → `yield* _(effect, trace)`
+ * - Modern pattern: `yield* effect` → `yield* $_adapter(effect, trace)`
+ *
+ * For the modern pattern, the transformer adds an adapter parameter to the generator
+ * function if not present, since Effect.gen always passes the adapter at runtime.
  *
  * @since 0.1.0
  */
 import type { NodePath, Visitor } from "@babel/traverse"
 import * as t from "@babel/types"
-import { isYieldAdapterCall } from "../utils/effectDetection.js"
+import {
+  getEffectGenGenerator,
+  isEffectGenCall,
+  isModernYield,
+  isYieldAdapterCallWithName
+} from "../utils/effectDetection.js"
 import {
   createHoistingState,
   extractLabelFromParent,
@@ -34,6 +43,30 @@ export interface SourceTraceOptions {
 interface TransformState {
   filename: string
   hoisting: HoistingState
+  /**
+   * Map from generator function to its adapter parameter name.
+   * Used to track which generators have been processed and their adapter names.
+   */
+  generatorAdapters: Map<unknown, string>
+}
+
+/**
+ * Default adapter name to use when adding one to modern pattern generators.
+ */
+const INJECTED_ADAPTER_NAME = "$"
+
+/**
+ * Gets the adapter parameter name from a generator function, or null if none.
+ */
+function getAdapterParamName(
+  gen: t.FunctionExpression | t.ArrowFunctionExpression
+): string | null {
+  if (gen.params.length === 0) return null
+  const firstParam = gen.params[0]
+  if (t.isIdentifier(firstParam)) {
+    return firstParam.name
+  }
+  return null
 }
 
 /**
@@ -48,6 +81,7 @@ export function createSourceTraceVisitor(
       enter(_path, state) {
         state.filename = filename
         state.hoisting = createHoistingState()
+        state.generatorAdapters = new Map()
       },
       exit(path, state) {
         // Prepend all hoisted statements to the program body
@@ -57,11 +91,45 @@ export function createSourceTraceVisitor(
       }
     },
 
+    // Process Effect.gen calls to detect and optionally add adapter parameters
+    CallExpression(path: NodePath<t.CallExpression>, state) {
+      const node = path.node
+
+      if (!isEffectGenCall(node)) return
+
+      const generator = getEffectGenGenerator(node)
+      if (!generator) return
+
+      // Check if generator already has an adapter parameter
+      const existingAdapter = getAdapterParamName(generator)
+
+      if (existingAdapter) {
+        // Generator uses legacy adapter pattern - store the name
+        state.generatorAdapters.set(generator, existingAdapter)
+      } else {
+        // Modern pattern - add adapter parameter
+        generator.params.unshift(t.identifier(INJECTED_ADAPTER_NAME))
+        state.generatorAdapters.set(generator, INJECTED_ADAPTER_NAME)
+      }
+    },
+
     YieldExpression(path: NodePath<t.YieldExpression>, state) {
       const node = path.node
 
-      // Only transform yield* _(effect) patterns
-      if (!isYieldAdapterCall(node)) {
+      // Must be yield* (delegating)
+      if (!node.delegate || !node.argument) return
+
+      // Find the enclosing generator function
+      let generatorPath = path.getFunctionParent()
+      if (!generatorPath) return
+
+      const generatorNode = generatorPath.node as t.FunctionExpression | t.ArrowFunctionExpression
+      if (!("generator" in generatorNode) || !generatorNode.generator) return
+
+      // Get the adapter name for this generator
+      const adapterName = state.generatorAdapters.get(generatorNode)
+      if (!adapterName) {
+        // Not inside an Effect.gen - skip
         return
       }
 
@@ -69,7 +137,7 @@ export function createSourceTraceVisitor(
       const loc = node.loc
       if (!loc) return
 
-      // Extract label from parent (e.g., `const user = yield* _(...)`)
+      // Extract label from parent (e.g., `const user = yield* ...`)
       const label = extractLabelFromParent(path)
 
       // Get or create hoisted trace identifier
@@ -81,9 +149,23 @@ export function createSourceTraceVisitor(
         label
       )
 
-      // Inject trace as second argument: _(effect) -> _(effect, trace)
-      const callExpr = node.argument as t.CallExpression
-      callExpr.arguments.push(t.cloneNode(traceId))
+      // Check if this is already an adapter call
+      if (
+        t.isCallExpression(node.argument) &&
+        isYieldAdapterCallWithName(node, adapterName)
+      ) {
+        // Legacy pattern: _(effect) -> _(effect, trace)
+        const callExpr = node.argument
+        callExpr.arguments.push(t.cloneNode(traceId))
+      } else if (isModernYield(node)) {
+        // Modern pattern: effect -> $(effect, trace)
+        const originalArg = node.argument
+        const wrappedCall = t.callExpression(
+          t.identifier(adapterName),
+          [originalArg, t.cloneNode(traceId)]
+        )
+        node.argument = wrappedCall
+      }
     }
   }
 }

package/src/transformers/withSpanTrace.ts

@@ -73,6 +73,62 @@ function isDataFirstCall(node: t.CallExpression): boolean {
   return false
 }
 
+/**
+ * Extracts the basename from a filepath.
+ */
+function basename(filepath: string): string {
+  const parts = filepath.split(/[/\\]/)
+  return parts[parts.length - 1] || filepath
+}
+
+/**
+ * Modifies a span name to include source location.
+ * "fetchUser" -> "fetchUser (index.ts:9)"
+ */
+function createSpanNameWithLocation(
+  originalName: string,
+  filepath: string,
+  line: number
+): string {
+  const filename = basename(filepath)
+  return `${originalName} (${filename}:${line})`
+}
+
+/**
+ * Creates the source location suffix string.
+ */
+function createLocationSuffix(filepath: string, line: number): string {
+  const filename = basename(filepath)
+  return ` (${filename}:${line})`
+}
+
+/**
+ * Modifies a template literal span name to include source location.
+ * Appends the location to the last quasi of the template.
+ */
+function modifyTemplateLiteralWithLocation(
+  templateLiteral: t.TemplateLiteral,
+  filepath: string,
+  line: number
+): t.TemplateLiteral {
+  const suffix = createLocationSuffix(filepath, line)
+  const quasis = [...templateLiteral.quasis]
+  const lastQuasi = quasis[quasis.length - 1]
+
+  // Create new quasi with appended location
+  const newLastQuasi = t.templateElement(
+    {
+      raw: lastQuasi.value.raw + suffix,
+      cooked: (lastQuasi.value.cooked || lastQuasi.value.raw) + suffix
+    },
+    lastQuasi.tail
+  )
+
+  quasis[quasis.length - 1] = newLastQuasi
+
+  return t.templateLiteral(quasis, [...templateLiteral.expressions])
+}
+
 /**
  * Creates the attributes object with source location.
  */
@@ -193,7 +249,15 @@ export function createWithSpanTraceVisitor(
 
       if (isDataFirst) {
         // Data-first: withSpan(effect, "name", options?)
-        // Options is at index 2
+        // Name is at index 1, options at index 2
+        const nameArg = node.arguments[1]
+        if (t.isStringLiteral(nameArg)) {
+          const newName = createSpanNameWithLocation(nameArg.value, state.filename, line)
+          node.arguments[1] = t.stringLiteral(newName)
+        } else if (t.isTemplateLiteral(nameArg)) {
+          node.arguments[1] = modifyTemplateLiteralWithLocation(nameArg, state.filename, line)
+        }
+
         const optionsArg = node.arguments[2]
         const newOptions = mergeOrCreateOptions(optionsArg, sourceAttrs)
 
@@ -206,7 +270,15 @@ export function createWithSpanTraceVisitor(
         }
       } else {
         // Data-last: withSpan("name", options?)
-        // Options is at index 1
+        // Name is at index 0, options at index 1
+        const nameArg = node.arguments[0]
+        if (t.isStringLiteral(nameArg)) {
+          const newName = createSpanNameWithLocation(nameArg.value, state.filename, line)
+          node.arguments[0] = t.stringLiteral(newName)
+        } else if (t.isTemplateLiteral(nameArg)) {
+          node.arguments[0] = modifyTemplateLiteralWithLocation(nameArg, state.filename, line)
+        }
+
         const optionsArg = node.arguments[1]
         const newOptions = mergeOrCreateOptions(optionsArg, sourceAttrs)
 

package/src/utils/effectDetection.ts

@@ -3,7 +3,7 @@
  *
  * @since 0.1.0
  */
-import type * as t from "@babel/types"
+import * as t from "@babel/types"
 
 /**
  * Checks if a CallExpression is `Effect.gen(...)`.
@@ -79,3 +79,51 @@ export function isYieldAdapterCall(node: t.YieldExpression): boolean {
   if (node.argument.type !== "CallExpression") return false
   return isAdapterCall(node.argument)
 }
+
+/**
+ * Checks if a YieldExpression is a `yield* _(effect)` pattern with a specific adapter name.
+ */
+export function isYieldAdapterCallWithName(node: t.YieldExpression, adapterName: string): boolean {
+  if (!node.delegate) return false // Must be yield*, not yield
+  if (!node.argument) return false
+  if (node.argument.type !== "CallExpression") return false
+  const callee = node.argument.callee
+  return callee.type === "Identifier" && callee.name === adapterName
+}
+
+/**
+ * Checks if a YieldExpression is a modern `yield* effect` pattern (without adapter).
+ * This is the pattern used when Effect.gen is called without the adapter parameter.
+ */
+export function isModernYield(node: t.YieldExpression): boolean {
+  if (!node.delegate) return false // Must be yield*, not yield
+  if (!node.argument) return false
+  // Modern yields are NOT adapter calls - they yield the effect directly
+  if (node.argument.type === "CallExpression" && isAdapterCall(node.argument)) {
+    return false
+  }
+  return true
+}
+
+/**
+ * Gets the generator function from an Effect.gen call.
+ * Returns the FunctionExpression/ArrowFunctionExpression if found.
+ */
+export function getEffectGenGenerator(node: t.CallExpression): t.FunctionExpression | t.ArrowFunctionExpression | null {
+  // Effect.gen can be called with 1 or 2 arguments:
+  // Effect.gen(function*() { ... }) - 1 arg
+  // Effect.gen(context, function*() { ... }) - 2 args
+  const args = node.arguments
+
+  for (let i = args.length - 1; i >= 0; i--) {
+    const arg = args[i]
+    if (t.isFunctionExpression(arg) && arg.generator) {
+      return arg
+    }
+    if (t.isArrowFunctionExpression(arg) && arg.generator) {
+      return arg
+    }
+  }
+
+  return null
+}