Skip to content

use effect concurrency #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nexxeln merged 1 commit into from
Dec 20, 2025
Merged

use effect concurrency #23

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
299 changes: 157 additions & 142 deletions packages/code-chunk/src/batch.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,12 @@
import { Effect, Queue, Stream } from 'effect'
import {
Effect,
Chunk as EffectChunk,
Exit,
Option,
Ref,
Scope,
Stream,
} from 'effect'
import { ChunkingError, UnsupportedLanguageError } from './chunk'
import { chunk as chunkInternal } from './chunking'
import { extractEntities } from './extract'
Expand All @@ -18,32 +26,156 @@ import type {

const DEFAULT_CONCURRENCY = 10

const chunkFileEffect = (
file: FileInput,
batchOptions: ChunkOptions = {},
): Effect.Effect<BatchResult, never> => {
const mergedOptions = { ...batchOptions, ...file.options }
/**
* Type for a function that chunks a single file and returns an Effect
*/
export type ChunkFileFunction = (
filepath: string,
code: string,
options: ChunkOptions,
) => Effect.Effect<Chunk[], unknown>

/**
* Core batch stream processor - takes a chunk function and returns a stream of results
* Used by both native and WASM implementations
*/
export const batchStreamEffect = (
chunkFile: ChunkFileFunction,
files: FileInput[],
options: BatchOptions = {},
): Stream.Stream<BatchResult, never> => {
const {
concurrency = DEFAULT_CONCURRENCY,
onProgress,
...chunkOptions
} = options
const total = files.length

if (total === 0) {
return Stream.empty
}

const processFile = (file: FileInput): Effect.Effect<BatchResult, never> => {
const mergedOptions = { ...chunkOptions, ...file.options }
return chunkFile(file.filepath, file.code, mergedOptions).pipe(
Effect.map(
(chunks) =>
({
filepath: file.filepath,
chunks,
error: null,
}) satisfies BatchFileResult,
),
Effect.catchAll((error) =>
Effect.succeed({
filepath: file.filepath,
chunks: null,
error: error instanceof Error ? error : new Error(String(error)),
} satisfies BatchFileError),
),
)
}

return Stream.unwrap(
Effect.gen(function* () {
const completedRef = yield* Ref.make(0)

return Stream.fromIterable(files).pipe(
Stream.mapEffect(
(file) =>
processFile(file).pipe(
Effect.tap((result) =>
Ref.updateAndGet(completedRef, (n) => n + 1).pipe(
Effect.andThen((completed) =>
Effect.sync(() =>
onProgress?.(
completed,
total,
file.filepath,
result.error === null,
),
),
),
),
),
),
{ concurrency },
),
)
}),
)
}

/**
* Core batch processor - collects stream results into array
*/
export const batchEffect = (
chunkFile: ChunkFileFunction,
files: FileInput[],
options: BatchOptions = {},
): Effect.Effect<BatchResult[], never> => {
return Stream.runCollect(batchStreamEffect(chunkFile, files, options)).pipe(
Effect.map((chunk) => Array.from(chunk)),
)
}

/**
* Core batch processor - Promise API
*/
export async function batch(
chunkFile: ChunkFileFunction,
files: FileInput[],
options?: BatchOptions,
): Promise<BatchResult[]> {
return Effect.runPromise(batchEffect(chunkFile, files, options))
}

/**
* Core batch stream processor - AsyncGenerator API
*/
export async function* batchStream(
chunkFile: ChunkFileFunction,
files: FileInput[],
options?: BatchOptions,
): AsyncGenerator<BatchResult> {
const scope = Effect.runSync(Scope.make())

try {
const pull = await Effect.runPromise(
Stream.toPull(batchStreamEffect(chunkFile, files, options)).pipe(
Scope.extend(scope),
),
)

while (true) {
const result = await Effect.runPromise(Effect.option(pull))
if (Option.isNone(result)) break
for (const item of EffectChunk.toReadonlyArray(result.value)) {
yield item
}
}
} finally {
await Effect.runPromise(Scope.close(scope, Exit.void))
}
}

const nativeChunkFile: ChunkFileFunction = (filepath, code, options) => {
return Effect.gen(function* () {
const language: Language | null =
mergedOptions.language ?? detectLanguage(file.filepath)
options.language ?? detectLanguage(filepath)

if (!language) {
return {
filepath: file.filepath,
chunks: null,
error: new UnsupportedLanguageError(file.filepath),
} satisfies BatchFileError
return yield* Effect.fail(new UnsupportedLanguageError(filepath))
}

const parseResult = yield* Effect.tryPromise({
try: () => parseCode(file.code, language),
try: () => parseCode(code, language),
catch: (error: unknown) =>
new ChunkingError('Failed to parse code', error),
})

const entities = yield* Effect.mapError(
extractEntities(parseResult.tree.rootNode, language, file.code),
extractEntities(parseResult.tree.rootNode, language, code),
(error: unknown) =>
new ChunkingError('Failed to extract entities', error),
)
Expand All @@ -57,163 +189,46 @@ const chunkFileEffect = (
const chunks = yield* Effect.mapError(
chunkInternal(
parseResult.tree.rootNode,
file.code,
code,
scopeTree,
language,
mergedOptions,
file.filepath,
options,
filepath,
),
(error: unknown) => new ChunkingError('Failed to chunk code', error),
)

const finalChunks: Chunk[] = parseResult.error
return parseResult.error
? chunks.map((c: Chunk) => ({
...c,
context: { ...c.context, parseError: parseResult.error ?? undefined },
}))
: chunks

return {
filepath: file.filepath,
chunks: finalChunks,
error: null,
} satisfies BatchFileResult
}).pipe(
Effect.catchAll((error) =>
Effect.succeed({
filepath: file.filepath,
chunks: null,
error: error instanceof Error ? error : new Error(String(error)),
} satisfies BatchFileError),
),
)
})
}

export const chunkBatchStreamEffect = (
files: FileInput[],
options: BatchOptions = {},
): Stream.Stream<BatchResult, never> => {
const {
concurrency = DEFAULT_CONCURRENCY,
onProgress,
...chunkOptions
} = options
const total = files.length

if (total === 0) {
return Stream.empty
}

return Stream.unwrap(
Effect.gen(function* () {
const queue = yield* Queue.unbounded<FileInput>()
const resultsQueue = yield* Queue.unbounded<BatchResult | null>()

yield* Effect.forEach(files, (file) => Queue.offer(queue, file), {
discard: true,
})

let completed = 0

const worker = Effect.gen(function* () {
while (true) {
const maybeFile = yield* Queue.poll(queue)
if (maybeFile._tag === 'None') {
break
}
const file = maybeFile.value
const result = yield* chunkFileEffect(file, chunkOptions)
completed++
if (onProgress) {
onProgress(completed, total, file.filepath, result.error === null)
}
yield* Queue.offer(resultsQueue, result)
}
})

yield* Effect.fork(
Effect.gen(function* () {
yield* Effect.all(
Array.from({ length: Math.min(concurrency, total) }, () => worker),
{ concurrency: 'unbounded' },
)
yield* Queue.offer(resultsQueue, null)
}),
)

return Stream.fromQueue(resultsQueue).pipe(
Stream.takeWhile((result): result is BatchResult => result !== null),
)
}),
)
}
): Stream.Stream<BatchResult, never> =>
batchStreamEffect(nativeChunkFile, files, options)

export const chunkBatchEffect = (
files: FileInput[],
options: BatchOptions = {},
): Effect.Effect<BatchResult[], never> => {
return Stream.runCollect(chunkBatchStreamEffect(files, options)).pipe(
Effect.map((chunk) => Array.from(chunk)),
)
}
): Effect.Effect<BatchResult[], never> =>
batchEffect(nativeChunkFile, files, options)

export async function chunkBatch(
files: FileInput[],
options?: BatchOptions,
): Promise<BatchResult[]> {
return Effect.runPromise(chunkBatchEffect(files, options))
return batch(nativeChunkFile, files, options)
}

export async function* chunkBatchStream(
files: FileInput[],
options?: BatchOptions,
): AsyncGenerator<BatchResult> {
const results: BatchResult[] = []
let resolveNext: ((value: IteratorResult<BatchResult>) => void) | null = null
let done = false

const streamEffect = chunkBatchStreamEffect(files, options).pipe(
Stream.runForEach((result) =>
Effect.sync(() => {
if (resolveNext) {
const resolve = resolveNext
resolveNext = null
resolve({ value: result, done: false })
} else {
results.push(result)
}
}),
),
Effect.tap(() =>
Effect.sync(() => {
done = true
if (resolveNext) {
resolveNext({ value: undefined as never, done: true })
}
}),
),
)

const runPromise = Effect.runPromise(streamEffect)

try {
while (true) {
const buffered = results.shift()
if (buffered !== undefined) {
yield buffered
} else if (done) {
break
} else {
const result = await new Promise<IteratorResult<BatchResult>>(
(resolve) => {
resolveNext = resolve
},
)
if (result.done) break
yield result.value
}
}
} finally {
await runPromise.catch(() => {})
}
yield* batchStream(nativeChunkFile, files, options)
}
Loading
Loading