refactor(core): refine layer node replacements (#34377)

packages/core/src/effect/app-node-builder.ts

@@ -1,24 +1,23 @@
-import { Layer } from "effect"
 import { buildLocationServiceMap } from "../location-services"
 import { LocationServiceMap } from "../location-service-map"
 import { LayerNode } from "./layer-node"
 import { makeGlobalNode } from "./app-node"
 
-export function build<A, E>(root: LayerNode.Node<A, E, any>, replacements?: readonly LayerNode.Replacement[]) {
-  const replacementMap = new Map(replacements?.map((item) => [item.source, item.replacement]))
+export function build<A, E>(root: LayerNode.Node<A, E, any>, replacements: LayerNode.Replacements = []) {
+  let allReplacements = replacements
 
-  if (!LayerNode.hasUnbound(root, LocationServiceMap.node)) {
-    // If the location service map is not needed, we shouldn't pull it
-    // in. Compile the graph normally
-    return LayerNode.compile(root, replacementMap)
+  // Only build the location service map if it's actually needed
+  if (LayerNode.hasUnbound(root, LocationServiceMap.node) && !hasReplacement(replacements, LocationServiceMap.node)) {
+    const locationMap = buildLocationServiceMap(replacements)
+    const locationMapNode = makeGlobalNode({ service: LocationServiceMap.Service, layer: locationMap, deps: [] })
+    allReplacements = replacements.concat([[LocationServiceMap.node, locationMapNode]])
   }
 
-  const locationMap = buildLocationServiceMap(replacementMap)
-  const locationMapNode = makeGlobalNode({ service: LocationServiceMap.Service, layer: locationMap, deps: [] })
+  return LayerNode.compile(root, allReplacements)
+}
 
-  const app = LayerNode.bind(root, LocationServiceMap.node, locationMapNode)
-
-  return LayerNode.compile(app, replacementMap)
+function hasReplacement(replacements: LayerNode.Replacements, node: LayerNode.Node<unknown, unknown, any>) {
+  return replacements.some(([source]) => source.name === node.name)
 }
 
 export * as AppNodeBuilder from "./app-node-builder"

packages/core/src/effect/dfdf

@@ -0,0 +1 @@
+File to save in: ~/.local/share/opencode/worktree/012780/location-layer-tiers/packages/core/src/effect/

packages/core/src/effect/layer-node.ts

@@ -111,20 +111,47 @@ export function group<const Items extends readonly AnyNode[]>(
   return { kind: "group", name: "group", dependencies }
 }
 
-export type Replacement = {
-  readonly source: Layer.Any
-  readonly replacement: Layer.Any
-}
+export type Replacement = readonly [source: AnyNode, replacement: AnyNode | Layer.Any]
+export type Replacements = readonly Replacement[]
 
 type CheckReplacementErrors<SourceError, ReplacementError> = [Exclude<ReplacementError, SourceError>] extends [never]
   ? unknown
   : { readonly "New replacement errors": Exclude<ReplacementError, SourceError> }
 
-export function replace<A, E, R, E2>(
-  source: Layer.Layer<A, E, R>,
-  replacement: Layer.Layer<NoInfer<A>, E2, never> & CheckReplacementErrors<E, NoInfer<E2>>,
-): Replacement {
-  return { source, replacement }
+type CheckReplacement<Item> = Item extends readonly [Node<infer A, infer E, infer T>, infer Replacement]
+  ? Replacement extends Node<NoInfer<A>, infer E2, T>
+    ? CheckReplacementErrors<E, NoInfer<E2>>
+    : Replacement extends Layer.Layer<NoInfer<A>, infer E2, never>
+      ? CheckReplacementErrors<E, NoInfer<E2>>
+      : { readonly "Invalid replacement": Replacement }
+  : { readonly "Invalid replacement": Item }
+
+type CheckReplacements<Items extends Replacements> = {
+  readonly [K in keyof Items]: CheckReplacement<Items[K]>
+}
+
+type ValidReplacements<Items extends Replacements> = Items & CheckReplacements<Items>
+
+function replacementNode(source: AnyNode, replacement: AnyNode | Layer.Any) {
+  const replacementNode = isNode(replacement)
+    ? replacement
+    : make({ ...nodeMakeIdentity(source), layer: replacement as Layer.Layer<unknown, unknown>, deps: [], tag: source.tag })
+  if (source.name !== replacementNode.name) {
+    throw new Error(`Cannot replace ${source.name} with ${replacementNode.name}`)
+  }
+  if (source.tag !== replacementNode.tag) {
+    throw new Error(`Cannot replace ${source.name} across tags`)
+  }
+  return replacementNode
+}
+
+function nodeMakeIdentity(node: AnyNode): NodeIdentity {
+  if (node.service !== undefined) return { service: node.service }
+  return { name: node.name }
+}
+
+function isNode(input: Layer.Any | AnyNode): input is AnyNode {
+  return "kind" in input && "dependencies" in input
 }
 
 // Tree -----------------------------------------------------------------------
@@ -176,32 +203,38 @@ function walk<Result>(
   return recur(root)
 }
 
-export function hoist<A, E, T extends Tag>(
+export function hoist<A, E, T extends Tag, const Items extends Replacements = readonly []>(
   root: Node<A, E, any>,
   tag: T,
+  replacements?: ValidReplacements<Items>,
 ): {
   readonly node: Node<A, E>
   readonly hoisted: Node<unknown, E>
 } {
   const hoisted = new Map<string, AnyNode>()
+  const replacementMap = replacementMapFrom(replacements)
 
-  const node = walk<AnyNode>(root, (node, context) => {
-    if (node.kind === "group") {
-      return { ...node, dependencies: node.dependencies.map(context.visit) }
-    }
-    if (node.tag === tag) {
-      const existing = hoisted.get(node.name)
-      if (existing && existing !== node) {
-        throw new Error(`Tag ${tag} has conflicting implementations for ${node.name}`)
+  const node = walk<AnyNode>(
+    root,
+    (node, context) => {
+      if (node.kind === "group") {
+        return { ...node, dependencies: node.dependencies.map(context.visit) }
       }
-      hoisted.set(node.name, node)
-      return group([])
-    }
-    if (node.kind === "unbound") {
-      return node
-    }
-    return { ...node, dependencies: node.dependencies.map(context.visit) }
-  })
+      if (node.tag === tag) {
+        const existing = hoisted.get(node.name)
+        if (existing && existing !== node) {
+          throw new Error(`Tag ${tag} has conflicting implementations for ${node.name}`)
+        }
+        hoisted.set(node.name, node)
+        return group([])
+      }
+      if (node.kind === "unbound") {
+        return node
+      }
+      return { ...node, dependencies: node.dependencies.map(context.visit) }
+    },
+    { resolve: (node) => replacementMap.get(node.name) ?? node },
+  )
 
   return {
     node: node as Node<A, E>,
@@ -209,10 +242,11 @@ export function hoist<A, E, T extends Tag>(
   }
 }
 
-export function compile<A, E>(
+export function compile<A, E, const Items extends Replacements = readonly []>(
   root: Node<A, E, any>,
-  replacements?: ReadonlyMap<Layer.Any, Layer.Any>,
+  replacements?: ValidReplacements<Items>,
 ): Layer.Layer<A, E> {
+  const replacementMap = replacementMapFrom(replacements)
   const cache = new Map<AnyNode, RuntimeLayer>()
   const compileNode = (node: AnyNode) =>
     walk<RuntimeLayer>(
@@ -220,18 +254,22 @@ export function compile<A, E>(
       (node, context) => {
         if (node.kind === "unbound") throw new Error(`Unbound layer node: ${node.name}`)
         const dependencies = node.dependencies.flatMap(flatten).map(context.visit)
-        const implementation = (replacements?.get(node.implementation!) ?? node.implementation!) as RuntimeLayer
+        const implementation = node.implementation! as RuntimeLayer
         return dependencies.length === 0
           ? implementation
           : implementation.pipe(Layer.provide(dependencies as [RuntimeLayer, ...RuntimeLayer[]]))
       },
-      { cache },
+      { cache, resolve: (node) => replacementMap.get(node.name) ?? node },
     )
   const layers = flatten(root).map((node) => compileNode(node))
   const layer = layers.reduce<RuntimeLayer>((result, layer) => layer.pipe(Layer.provideMerge(result)), Layer.empty)
   return layer as Layer.Layer<A, E>
 }
 
+function replacementMapFrom(replacements?: Replacements) {
+  return new Map(replacements?.map(([source, replacement]) => [source.name, replacementNode(source, replacement)]))
+}
+
 export function hasUnbound(root: Node<unknown, unknown, any>, source: AnyNode): boolean {
   if (source.kind !== "unbound") throw new Error(`Cannot check non-unbound layer node: ${source.name}`)
   return walk<boolean>(root, (node, context) => {
@@ -240,32 +278,6 @@ export function hasUnbound(root: Node<unknown, unknown, any>, source: AnyNode):
   })
 }
 
-export function bind<A, E, T extends Tag | undefined>(
-  root: Node<A, E, T>,
-  source: AnyNode,
-  replacement: AnyNode,
-): Node<A, E, T> {
-  if (source.kind !== "unbound") throw new Error(`Cannot bind non-unbound layer node: ${source.name}`)
-  if (source.name !== replacement.name) {
-    throw new Error(`Cannot bind ${source.name} to ${replacement.name}`)
-  }
-  if (source.tag !== replacement.tag) {
-    throw new Error(`Cannot bind ${source.name} across tags`)
-  }
-  return walk<AnyNode>(
-    root,
-    (target, context) => {
-      if (target.kind === "unbound") return target
-      const dependencies: AnyNode[] = []
-      const clone = { ...target, dependencies }
-      context.cache.set(target, clone)
-      dependencies.push(...target.dependencies.map(context.visit))
-      return clone
-    },
-    { detectCycles: false, resolve: (node) => (node === source ? replacement : node) },
-  ) as Node<A, E, T>
-}
-
 function flatten(node: AnyNode): readonly AnyNode[] {
   return node.kind === "group" ? node.dependencies.flatMap(flatten) : [node]
 }

packages/core/src/location-services.ts

@@ -82,17 +82,16 @@ export type LocationServices = LayerNode.Output<typeof locationServices>
 export type LocationError = LayerNode.Error<typeof locationServices>
 
 export function buildLocationServiceMap(
-  replacements?: ReadonlyMap<Layer.Any, Layer.Any>,
+  replacements: LayerNode.Replacements = [],
 ): Layer.Layer<LocationServiceMap.Service> {
   return Layer.effect(
     LocationServiceMap.Service,
     LayerMap.make(
       (ref: Location.Ref) => {
-        const location = LayerNode.hoist(
-          LayerNode.bind(locationServices, Location.node, Location.boundNode(ref)),
-          Node.tags.values.global,
-        )
-        return LayerNode.compile(location.node, replacements).pipe(
+        const allReplacements = replacements.concat([[Location.node, Location.boundNode(ref)]])
+        const location = LayerNode.hoist(locationServices, Node.tags.values.global, allReplacements)
+
+        return LayerNode.compile(location.node).pipe(
           Layer.fresh,
           Layer.tap(() =>
             Effect.logInfo("booting location services", {
@@ -100,7 +99,7 @@ export function buildLocationServiceMap(
               workspaceID: ref.workspaceID,
             }),
           ),
-          Layer.provide(LayerNode.compile(location.hoisted, replacements)),
+          Layer.provide(LayerNode.compile(location.hoisted)),
         )
       },
       { idleTimeToLive: "60 minutes" },

packages/core/src/tool-output-store.ts

@@ -196,6 +196,8 @@ export const defaultLayer = layer.pipe(Layer.provide(FSUtil.defaultLayer), Layer
 
 export const node = makeLocationNode({ service: Service, layer, deps: [FSUtil.node, Global.node, Config.node] })
 
+export const nodeWithoutConfig = makeLocationNode({ service: Service, layer, deps: [FSUtil.node, Global.node] })
+
 /** Runs retention scanning once globally rather than once per active Location. */
 export const cleanupLayer = Layer.effectDiscard(
   Effect.gen(function* () {