Add reverse implementors index to the partial SCIP loader

src/scip-lib/partialloader/index.go

@@ -24,6 +24,7 @@ type PartialIndex interface {
 	GetSymbolInformation(symbol string) (*model.SymbolInformation, string, error)
 	GetSymbolInformationFromDescriptors(descriptors []model.Descriptor, version string) (*model.SymbolInformation, string, error)
 	References(symbol string) (map[string][]*model.Occurrence, error)
+	Implementations(symbol string) ([]string, error)
 	Tidy() error
 }
 
@@ -57,19 +58,24 @@ type PartialLoadedIndex struct {
 	indexFolder      string
 	pool             *scanner.BufferPool
 	onDocumentLoaded func(*model.Document)
+
+	// ImplementorsBySymbol maps abstract/interface symbol -> set of implementing symbols
+	implementorsMu       sync.RWMutex
+	ImplementorsBySymbol map[string]map[string]struct{}
 }
 
 // NewPartialLoadedIndex creates a new PartialLoadedIndex
 func NewPartialLoadedIndex(indexFolder string) PartialIndex {
 	return &PartialLoadedIndex{
-		PrefixTreeRoot:   NewSymbolPrefixTree(),
-		DocTreeNodes:     make(map[string]*docNodes),
-		LoadedDocuments:  make(map[string]*model.Document),
-		updatedDocs:      make(map[string]int64),
-		docToIndex:       make(map[string]string, 0),
-		indexFolder:      indexFolder,
-		pool:             scanner.NewBufferPool(1024, 12),
-		onDocumentLoaded: func(*model.Document) {},
+		PrefixTreeRoot:       NewSymbolPrefixTree(),
+		DocTreeNodes:         make(map[string]*docNodes),
+		LoadedDocuments:      make(map[string]*model.Document),
+		updatedDocs:          make(map[string]int64),
+		docToIndex:           make(map[string]string, 0),
+		indexFolder:          indexFolder,
+		pool:                 scanner.NewBufferPool(1024, 12),
+		onDocumentLoaded:     func(*model.Document) {},
+		ImplementorsBySymbol: make(map[string]map[string]struct{}),
 	}
 }
 
@@ -130,6 +136,7 @@ func (p *PartialLoadedIndex) LoadIndex(indexPath string, indexReader scanner.Sci
 	localDocTreeNodes := make(map[string]*docNodes)
 	localUpdatedDocs := make(map[string]int64)
 	localDocToIndex := make(map[string]string)
+	localImplementorsBySymbol := make(map[string]map[string]struct{})
 
 	loadScanner := &scanner.IndexScannerImpl{
 		Pool: p.pool,
@@ -157,6 +164,17 @@ func (p *PartialLoadedIndex) LoadIndex(indexPath string, indexReader scanner.Sci
 			modelInfo := mapper.ScipSymbolInformationToModelSymbolInformation(info)
 			leafNode, isNew := localPrefixTree.AddSymbol(docPath, modelInfo, p.revision.Load())
 
+			// Populate reverse implementors for quick lookup (impl -> abs)
+			for _, rel := range modelInfo.Relationships {
+				if rel != nil && rel.IsImplementation {
+					abs := rel.Symbol
+					if localImplementorsBySymbol[abs] == nil {
+						localImplementorsBySymbol[abs] = make(map[string]struct{})
+					}
+					localImplementorsBySymbol[abs][modelInfo.Symbol] = struct{}{}
+				}
+			}
+
 			if isNew {
 				localDocTreeNodes[docPath].nodes = append(localDocTreeNodes[docPath].nodes, leafNode)
 			}
@@ -179,6 +197,7 @@ func (p *PartialLoadedIndex) LoadIndex(indexPath string, indexReader scanner.Sci
 		p.mergeDocTreeNodes(localDocTreeNodes)
 		p.mergeUpdatedDocs(localUpdatedDocs)
 		p.mergeDocToIndex(localDocToIndex)
+		p.mergeImplementors(localImplementorsBySymbol)
 	}()
 	return loadScanner.ScanIndexReader(indexReader)
 }
@@ -280,6 +299,23 @@ func (p *PartialLoadedIndex) mergeDocToIndex(localDocToIndex map[string]string)
 	}
 }
 
+// mergeImplementors merges a local reverse implementors map into the main index
+func (p *PartialLoadedIndex) mergeImplementors(local map[string]map[string]struct{}) {
+	if len(local) == 0 {
+		return
+	}
+	p.implementorsMu.Lock()
+	defer p.implementorsMu.Unlock()
+	for abs, impls := range local {
+		if p.ImplementorsBySymbol[abs] == nil {
+			p.ImplementorsBySymbol[abs] = make(map[string]struct{})
+		}
+		for impl := range impls {
+			p.ImplementorsBySymbol[abs][impl] = struct{}{}
+		}
+	}
+}
+
 // LoadDocument loads a document into the PartialLoadedIndex
 func (p *PartialLoadedIndex) LoadDocument(relativeDocPath string) (*model.Document, error) {
 	p.loadedDocsMu.RLock()
@@ -364,6 +400,21 @@ func (p *PartialLoadedIndex) loadDocumentFromIndexFolder(relativeDocPath string)
 	return doc, nil
 }
 
+// Implementations returns the list of implementing symbols for a given abstract/interface symbol
+func (p *PartialLoadedIndex) Implementations(symbol string) ([]string, error) {
+	p.implementorsMu.RLock()
+	defer p.implementorsMu.RUnlock()
+	set := p.ImplementorsBySymbol[symbol]
+	if set == nil {
+		return []string{}, nil
+	}
+	res := make([]string, 0, len(set))
+	for s := range set {
+		res = append(res, s)
+	}
+	return res, nil
+}
+
 // Tidy prunes nodes for documents that were updated in the current revision
 func (p *PartialLoadedIndex) Tidy() error {
 	// Acquire the modification mutex to prevent new index loads during cleanup

src/scip-lib/partialloader/index_test.go

@@ -410,3 +410,52 @@ func TestLoadIndexWithPreloadedDocument(t *testing.T) {
 	assert.NotNil(t, info)
 	assert.Equal(t, docPath, foundDocPath)
 }
+
+// New tests for reverse implementors index
+func TestMergeImplementors(t *testing.T) {
+	idx := &PartialLoadedIndex{
+		ImplementorsBySymbol: make(map[string]map[string]struct{}),
+	}
+	local := map[string]map[string]struct{}{
+		"abs#Symbol": {
+			"impl#A": {},
+			"impl#B": {},
+		},
+	}
+	idx.mergeImplementors(local)
+	set := idx.ImplementorsBySymbol["abs#Symbol"]
+	if assert.NotNil(t, set) {
+		_, okA := set["impl#A"]
+		_, okB := set["impl#B"]
+		assert.True(t, okA)
+		assert.True(t, okB)
+	}
+}
+
+func TestImplementations(t *testing.T) {
+	idx := &PartialLoadedIndex{
+		ImplementorsBySymbol: make(map[string]map[string]struct{}),
+	}
+	idx.ImplementorsBySymbol["abs#Symbol"] = map[string]struct{}{
+		"impl#B": {},
+		"impl#A": {},
+	}
+	list, err := idx.Implementations("abs#Symbol")
+	assert.NoError(t, err)
+	assert.Contains(t, list, "impl#A")
+	assert.Contains(t, list, "impl#B")
+
+	empty, err := idx.Implementations("unknown#Symbol")
+	assert.NoError(t, err)
+	assert.Equal(t, 0, len(empty))
+}
+
+func TestLoadIndexWithImplementors(t *testing.T) {
+	index := NewPartialLoadedIndex("../testdata")
+	err := index.LoadIndexFile(filepath.Join("../testdata", "index.scip"))
+	assert.NoError(t, err)
+	symbol := "scip-go gomod code.uber.internal/devexp/test_management/tracing 0f67d80e60274b77875a241c43ef980bc9ffe0d8 `code.uber.internal/devexp/test_management/tracing`/PartialIndex#"
+	implementors, err := index.Implementations(symbol)
+	assert.NoError(t, err)
+	assert.Equal(t, []string{"scip-go gomod code.uber.internal/devexp/test_management/tracing 0f67d80e60274b77875a241c43ef980bc9ffe0d8 `code.uber.internal/devexp/test_management/tracing`/index#"}, implementors)
+}

src/ulsp/controller/scip/partial_registry.go

@@ -319,6 +319,56 @@ func (p *partialScipRegistry) GetSymbolOccurrence(uri uri.URI, pos protocol.Posi
 	return nil, errors.New("not implemented")
 }
 
+// Implementation implements Registry.
+func (p *partialScipRegistry) Implementation(sourceURI uri.URI, pos protocol.Position) ([]protocol.Location, error) {
+	doc, err := p.Index.LoadDocument(p.uriToRelativePath(sourceURI))
+	if err != nil {
+		p.logger.Errorf("failed to load document %s: %s", sourceURI, err)
+		return nil, err
+	}
+	if doc == nil {
+		return nil, nil
+	}
+
+	sourceOccurrence := GetOccurrenceForPosition(doc.Occurrences, pos)
+	if sourceOccurrence == nil {
+		return nil, nil
+	}
+
+	// Local symbols typically don't have implementation relationships
+	if scip.IsLocalSymbol(sourceOccurrence.Symbol) {
+		return []protocol.Location{}, nil
+	}
+
+	locations := make([]protocol.Location, 0)
+
+	implementors, err := p.Index.Implementations(sourceOccurrence.Symbol)
+	if err != nil {
+		p.logger.Errorf("failed to get implementing symbols for %s: %s", sourceOccurrence.Symbol, err)
+		return nil, err
+	}
+	for _, implSym := range implementors {
+		implementingSymbol, err := model.ParseScipSymbol(implSym)
+		if err != nil {
+			p.logger.Errorf("failed to parse implementing symbol %s: %s", implSym, err)
+			continue
+		}
+		implOcc, err := p.GetSymbolDefinitionOccurrence(
+			mapper.ScipDescriptorsToModelDescriptors(implementingSymbol.Descriptors),
+			implementingSymbol.Package.Version,
+		)
+		if err != nil {
+			p.logger.Errorf("failed to get definition for implementing symbol %s: %s", implSym, err)
+			continue
+		}
+		if implOcc != nil && implOcc.Occurrence != nil {
+			locations = append(locations, *mapper.ScipOccurrenceToLocation(implOcc.Location, implOcc.Occurrence))
+		}
+	}
+
+	return locations, nil
+}
+
 func (p *partialScipRegistry) uriToRelativePath(uri uri.URI) string {
 	rel, err := filepath.Rel(p.WorkspaceRoot, uri.Filename())
 	if err != nil {

src/ulsp/controller/scip/partial_registry_test.go

@@ -1475,3 +1475,41 @@ func TestPartialScipRegistry_DocumentSymbols(t *testing.T) {
 		})
 	}
 }
+
+func TestPartialScipRegistry_Implementation_FastPath(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mockIndex := partialloadermock.NewMockPartialIndex(ctrl)
+	logger := zaptest.NewLogger(t).Sugar()
+
+	registry := &partialScipRegistry{
+		WorkspaceRoot: "/workspace",
+		Index:         mockIndex,
+		logger:        logger,
+	}
+
+	sourceURI := uri.File("/workspace/test.go")
+	pos := protocol.Position{Line: 1, Character: 1}
+
+	sourceOcc := &model.Occurrence{Symbol: tracingUUIDKey, Range: []int32{1, 1, 1, 2}}
+	mockIndex.EXPECT().LoadDocument("test.go").Return(&model.Document{
+		Occurrences: []*model.Occurrence{sourceOcc},
+	}, nil)
+
+	mockIndex.EXPECT().Implementations(tracingUUIDKey).Return([]string{
+		"scip-go gomod example/pkg v1.0.0 `example/pkg`/Foo#Bar.",
+	}, nil)
+
+	mockIndex.EXPECT().GetSymbolInformationFromDescriptors(gomock.Any(), gomock.Any()).Return(&model.SymbolInformation{Symbol: "impl#sym"}, "impl.go", nil)
+	mockIndex.EXPECT().LoadDocument("impl.go").Return(&model.Document{
+		Occurrences: []*model.Occurrence{
+			{Symbol: "impl#sym", SymbolRoles: int32(scipproto.SymbolRole_Definition), Range: []int32{10, 1, 10, 5}},
+		},
+	}, nil)
+
+	locs, err := registry.Implementation(sourceURI, pos)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(locs))
+	assert.Equal(t, uri.File("/workspace/impl.go"), locs[0].URI)
+	assert.Equal(t, protocol.Position{Line: 10, Character: 1}, locs[0].Range.Start)
+	assert.Equal(t, protocol.Position{Line: 10, Character: 5}, locs[0].Range.End)
+}