diff --git a/src/context/index.ts b/src/context/index.ts index b6efd18..7b666ed 100644 --- a/src/context/index.ts +++ b/src/context/index.ts @@ -5,10 +5,12 @@ import type { ByteRange, Chunk, ChunkContext, + ChunkEntityInfo, ChunkOptions, EntityInfo, ExtractedEntity, ImportInfo, + Language, ScopeTree, } from '../types' import { getSiblings, type SiblingOptions } from './siblings' @@ -68,6 +70,26 @@ export const getScopeForRange = ( return scopeChain } +/** + * Check if an entity is partial (not fully contained) within a byte range + * + * An entity is partial if it overlaps with the range but is not fully contained. + * + * @param entity - The entity to check + * @param byteRange - The chunk's byte range + * @returns Whether the entity is partial + */ +const isEntityPartial = ( + entity: ExtractedEntity, + byteRange: ByteRange, +): boolean => { + // Entity is partial if it starts before the range or ends after the range + return ( + entity.byteRange.start < byteRange.start || + entity.byteRange.end > byteRange.end + ) +} + /** * Get entities within a byte range * @@ -76,7 +98,7 @@ export const getScopeForRange = ( * * @param byteRange - The byte range to search * @param scopeTree - The scope tree - * @returns Entity info array for entities in range + * @returns Entity info array for entities in range with isPartial detection */ export const getEntitiesInRange = ( byteRange: ByteRange, @@ -90,12 +112,17 @@ export const getEntitiesInRange = ( ) }) - // Map to EntityInfo - return overlappingEntities.map((entity) => ({ - name: entity.name, - type: entity.type, - signature: entity.signature, - })) + // Map to ChunkEntityInfo with additional fields + return overlappingEntities.map( + (entity): ChunkEntityInfo => ({ + name: entity.name, + type: entity.type, + signature: entity.signature, + docstring: entity.docstring, + lineRange: entity.lineRange, + isPartial: isEntityPartial(entity, byteRange), + }), + ) } /** @@ -169,23 +196,38 @@ export const getRelevantImports = ( return filteredImports.map(mapToImportInfo) } +/** + * Options for attaching context to a chunk + */ +export interface AttachContextOptions { + /** The rebuilt text info for the chunk */ + text: RebuiltText + /** The scope tree for the file */ + scopeTree: ScopeTree + /** Chunking options */ + options: ChunkOptions + /** The chunk index */ + index: number + /** Total number of chunks */ + totalChunks: number + /** File path of the source file (optional) */ + filepath?: string + /** Programming language of the source (optional) */ + language?: Language +} + /** * Attach context information to a chunk * - * @param text - The rebuilt text info for the chunk - * @param scopeTree - The scope tree for the file - * @param options - Chunking options - * @param index - The chunk index - * @param totalChunks - Total number of chunks + * @param opts - Options containing all parameters for context attachment * @returns Effect yielding the complete chunk with context */ export const attachContext = ( - text: RebuiltText, - scopeTree: ScopeTree, - options: ChunkOptions, - index: number, - totalChunks: number, + opts: AttachContextOptions, ): Effect.Effect => { + const { text, scopeTree, options, index, totalChunks, filepath, language } = + opts + return Effect.try({ try: () => { // Determine context mode @@ -194,6 +236,8 @@ export const attachContext = ( // For 'none' mode, return minimal context if (contextMode === 'none') { const context: ChunkContext = { + filepath, + language, scope: [], entities: [], siblings: [], @@ -228,6 +272,8 @@ export const attachContext = ( const imports = getRelevantImports(entities, scopeTree, filterImports) const context: ChunkContext = { + filepath, + language, scope, entities, siblings, diff --git a/src/types.ts b/src/types.ts index 822f05c..28269b4 100644 --- a/src/types.ts +++ b/src/types.ts @@ -158,6 +158,19 @@ export interface EntityInfo { signature?: string } +/** + * Extended entity info for entities within a chunk + * Includes additional context like docstring, line range, and partial status + */ +export interface ChunkEntityInfo extends EntityInfo { + /** Documentation comment if present */ + docstring?: string | null + /** Line range in source (0-indexed, inclusive) */ + lineRange?: LineRange + /** Whether this entity spans multiple chunks (is partial) */ + isPartial?: boolean +} + /** * Information about a sibling entity */ @@ -190,10 +203,14 @@ export interface ImportInfo { * Context information for a chunk */ export interface ChunkContext { - /** Scope information */ + /** File path of the source file */ + filepath?: string + /** Programming language of the source */ + language?: Language + /** Scope information (scope chain from current to root) */ scope: EntityInfo[] /** Entities within this chunk */ - entities: EntityInfo[] + entities: ChunkEntityInfo[] /** Nearby sibling entities */ siblings: SiblingInfo[] /** Relevant imports */ diff --git a/test/scope.test.ts b/test/scope.test.ts index 8ad9c7f..43dffb1 100644 --- a/test/scope.test.ts +++ b/test/scope.test.ts @@ -468,3 +468,155 @@ func subtract(a, b int) int { expect(cls).toBeDefined() }) }) + +// ============================================================================ +// Context Attachment Tests +// ============================================================================ + +describe('context attachment', () => { + test('getEntitiesInRange returns entities with isPartial flag', async () => { + const code = `function foo() { return 1 } +function bar() { return 2 } +function baz() { return 3 }` + const entities = await getEntities(code, 'typescript') + const tree = buildScopeTreeFromEntities(entities) + + // Import the function we need to test + const { getEntitiesInRange } = await import('../src/context/index') + + // Get entities for a range that fully contains 'bar' but not 'foo' or 'baz' + const barEntity = entities.find((e) => e.name === 'bar') + if (barEntity) { + const entitiesInRange = getEntitiesInRange(barEntity.byteRange, tree) + + // Should find bar + const bar = entitiesInRange.find((e) => e.name === 'bar') + expect(bar).toBeDefined() + // bar should NOT be partial since we're using its exact range + expect(bar?.isPartial).toBe(false) + } + }) + + test('getEntitiesInRange marks partial entities correctly', async () => { + const code = `class BigClass { + method1() { return 1 } + method2() { return 2 } + method3() { return 3 } +}` + const entities = await getEntities(code, 'typescript') + const tree = buildScopeTreeFromEntities(entities) + + const { getEntitiesInRange } = await import('../src/context/index') + + // Get just method2's range - this should be inside BigClass + const method2 = entities.find((e) => e.name === 'method2') + if (method2) { + const entitiesInRange = getEntitiesInRange(method2.byteRange, tree) + + // method2 should not be partial (its full range is included) + const m2 = entitiesInRange.find((e) => e.name === 'method2') + expect(m2?.isPartial).toBe(false) + + // BigClass should be partial (we only have a slice of it) + const cls = entitiesInRange.find((e) => e.name === 'BigClass') + if (cls) { + expect(cls.isPartial).toBe(true) + } + } + }) + + test('getEntitiesInRange includes docstring and lineRange', async () => { + const code = `/** + * A test function with docs. + */ +function documented() { + return 1 +}` + const entities = await getEntities(code, 'typescript') + const tree = buildScopeTreeFromEntities(entities) + + const { getEntitiesInRange } = await import('../src/context/index') + + const fn = entities.find((e) => e.name === 'documented') + if (fn) { + const entitiesInRange = getEntitiesInRange(fn.byteRange, tree) + const docFn = entitiesInRange.find((e) => e.name === 'documented') + + expect(docFn).toBeDefined() + expect(docFn?.lineRange).toBeDefined() + // Docstring should be present if extracted + if (fn.docstring) { + expect(docFn?.docstring).toContain('test function') + } + } + }) + + test('attachContext includes filepath and language', async () => { + const { Effect } = await import('effect') + const { attachContext } = await import('../src/context/index') + + const code = `function test() { return 1 }` + const entities = await getEntities(code, 'typescript') + const tree = buildScopeTreeFromEntities(entities) + + const fn = entities[0] + if (fn) { + const mockText = { + text: code, + byteRange: { start: 0, end: code.length }, + lineRange: { start: 0, end: 0 }, + } + + const chunk = await Effect.runPromise( + attachContext({ + text: mockText, + scopeTree: tree, + options: {}, + index: 0, + totalChunks: 1, + filepath: 'test.ts', + language: 'typescript', + }), + ) + + expect(chunk.context.filepath).toBe('test.ts') + expect(chunk.context.language).toBe('typescript') + } + }) + + test('attachContext respects contextMode none', async () => { + const { Effect } = await import('effect') + const { attachContext } = await import('../src/context/index') + + const code = `function test() { return 1 }` + const entities = await getEntities(code, 'typescript') + const tree = buildScopeTreeFromEntities(entities) + + const mockText = { + text: code, + byteRange: { start: 0, end: code.length }, + lineRange: { start: 0, end: 0 }, + } + + const chunk = await Effect.runPromise( + attachContext({ + text: mockText, + scopeTree: tree, + options: { contextMode: 'none' }, + index: 0, + totalChunks: 1, + filepath: 'test.ts', + language: 'typescript', + }), + ) + + // Even in 'none' mode, filepath and language should be present + expect(chunk.context.filepath).toBe('test.ts') + expect(chunk.context.language).toBe('typescript') + // But scope, entities, siblings, imports should be empty + expect(chunk.context.scope).toEqual([]) + expect(chunk.context.entities).toEqual([]) + expect(chunk.context.siblings).toEqual([]) + expect(chunk.context.imports).toEqual([]) + }) +})