zzet--gortex/internal/parser/languages/ocaml.go

package languages

import (
	"fmt"
	"strings"

	"github.com/zzet/gortex/internal/graph"
	"github.com/zzet/gortex/internal/parser"
	sitter "github.com/zzet/gortex/internal/parser/tsitter"
	"github.com/zzet/gortex/internal/parser/tsitter/ocaml"
)

// OCamlExtractor extracts OCaml source files.
type OCamlExtractor struct {
	lang *sitter.Language
}

func NewOCamlExtractor() *OCamlExtractor {
	return &OCamlExtractor{lang: ocaml.GetLanguage()}
}

func (e *OCamlExtractor) Language() string     { return "ocaml" }
func (e *OCamlExtractor) Extensions() []string { return []string{".ml", ".mli"} }

func (e *OCamlExtractor) Extract(filePath string, src []byte) (*parser.ExtractionResult, error) {
	tree, err := parser.ParseFile(src, e.lang)
	if err != nil {
		return nil, err
	}
	defer tree.Close()

	root := tree.RootNode()
	result := &parser.ExtractionResult{}

	fileNode := &graph.Node{
		ID: filePath, Kind: graph.KindFile, Name: filePath,
		FilePath: filePath, StartLine: 1, EndLine: int(root.EndPoint().Row) + 1,
		Language: "ocaml",
	}
	result.Nodes = append(result.Nodes, fileNode)

	seen := make(map[string]bool)

	// Walk the AST to extract definitions.
	e.walkExtract(root, src, filePath, fileNode, result, seen, "")

	// Call sites.
	funcRanges := buildFuncRanges(result)
	e.extractCalls(root, src, filePath, result, funcRanges)

	return result, nil
}

func (e *OCamlExtractor) walkExtract(
	node *sitter.Node, src []byte, filePath string, fileNode *graph.Node,
	result *parser.ExtractionResult, seen map[string]bool, modulePrefix string,
) {
	for i, _nc := 0, int(node.ChildCount()); i < _nc; i++ {
		child := node.Child(i)
		if child == nil {
			continue
		}

		switch child.Type() {
		case "value_definition":
			e.extractValueDef(child, src, filePath, fileNode, result, seen, modulePrefix)

		case "type_definition":
			e.extractTypeDef(child, src, filePath, fileNode, result, seen, modulePrefix)

		case "module_definition":
			e.extractModuleDef(child, src, filePath, fileNode, result, seen, modulePrefix)

		case "module_type_definition":
			e.extractModuleTypeDef(child, src, filePath, fileNode, result, seen, modulePrefix)

		case "open_module":
			e.extractOpen(child, src, filePath, fileNode, result)

		case "class_definition":
			e.extractClassDef(child, src, filePath, fileNode, result, seen, modulePrefix)

		case "external":
			e.extractExternal(child, src, filePath, fileNode, result, seen, modulePrefix)

		case "value_specification":
			// .mli signature files
			e.extractValueSpec(child, src, filePath, fileNode, result, seen, modulePrefix)
		}
	}
}

// extractValueDef handles `let name = ...` and `let rec name ... = ...`
func (e *OCamlExtractor) extractValueDef(
	node *sitter.Node, src []byte, filePath string, fileNode *graph.Node,
	result *parser.ExtractionResult, seen map[string]bool, modulePrefix string,
) {
	for i, _nc := 0, int(node.NamedChildCount()); i < _nc; i++ {
		child := node.NamedChild(i)
		if child == nil {
			continue
		}

		if child.Type() == "let_binding" {
			name := ""
			kind := graph.KindFunction
			hasParams := false

			// Walk let_binding children to find pattern (name) and check for parameters.
			for j, _nc := 0, int(child.NamedChildCount()); j < _nc; j++ {
				part := child.NamedChild(j)
				if part == nil {
					continue
				}
				switch part.Type() {
				case "value_name", "value_pattern":
					name = part.Content(src)
				case "parameter":
					hasParams = true
				case "fun_expression", "function_expression":
					hasParams = true
				}
			}

			if name == "" || name == "_" {
				continue
			}

			if !hasParams {
				kind = graph.KindVariable
			}

			qualName := name
			if modulePrefix != "" {
				qualName = modulePrefix + "." + name
			}

			id := filePath + "::" + qualName
			if seen[id] {
				continue
			}
			seen[id] = true

			startLine := int(child.StartPoint().Row) + 1
			endLine := int(child.EndPoint().Row) + 1

			n := &graph.Node{
				ID: id, Kind: kind, Name: name,
				FilePath: filePath, StartLine: startLine, EndLine: endLine,
				Language: "ocaml",
			}
			if hasParams {
				n.Meta = map[string]any{"signature": "let " + name + " ..."}
			}
			result.Nodes = append(result.Nodes, n)
			result.Edges = append(result.Edges, &graph.Edge{
				From: fileNode.ID, To: id, Kind: graph.EdgeDefines,
				FilePath: filePath, Line: startLine,
			})
		}
	}
}

// extractTypeDef handles `type name = ...`
func (e *OCamlExtractor) extractTypeDef(
	node *sitter.Node, src []byte, filePath string, fileNode *graph.Node,
	result *parser.ExtractionResult, seen map[string]bool, modulePrefix string,
) {
	for i, _nc := 0, int(node.NamedChildCount()); i < _nc; i++ {
		child := node.NamedChild(i)
		if child == nil || child.Type() != "type_binding" {
			continue
		}

		name := ""
		for j, _nc := 0, int(child.NamedChildCount()); j < _nc; j++ {
			part := child.NamedChild(j)
			if part != nil && (part.Type() == "type_constructor" || part.Type() == "type_variable") {
				name = part.Content(src)
				break
			}
		}

		if name == "" {
			// Try first named child as name.
			if child.NamedChildCount() > 0 {
				name = child.NamedChild(0).Content(src)
			}
		}

		if name == "" {
			continue
		}

		qualName := name
		if modulePrefix != "" {
			qualName = modulePrefix + "." + name
		}

		id := filePath + "::" + qualName
		if seen[id] {
			continue
		}
		seen[id] = true

		startLine := int(child.StartPoint().Row) + 1
		endLine := int(child.EndPoint().Row) + 1

		result.Nodes = append(result.Nodes, &graph.Node{
			ID: id, Kind: graph.KindType, Name: name,
			FilePath: filePath, StartLine: startLine, EndLine: endLine,
			Language: "ocaml", Meta: map[string]any{"type_flavor": "type_def"},
		})
		result.Edges = append(result.Edges, &graph.Edge{
			From: fileNode.ID, To: id, Kind: graph.EdgeDefines,
			FilePath: filePath, Line: startLine,
		})
	}
}

// extractModuleDef handles `module Name = struct ... end`
// AST: module_definition → module_binding → module_name + structure
func (e *OCamlExtractor) extractModuleDef(
	node *sitter.Node, src []byte, filePath string, fileNode *graph.Node,
	result *parser.ExtractionResult, seen map[string]bool, modulePrefix string,
) {
	// Find module_binding child.
	for i, _nc := 0, int(node.NamedChildCount()); i < _nc; i++ {
		binding := node.NamedChild(i)
		if binding == nil || binding.Type() != "module_binding" {
			continue
		}

		name := ""
		for j, _nc := 0, int(binding.NamedChildCount()); j < _nc; j++ {
			child := binding.NamedChild(j)
			if child != nil && child.Type() == "module_name" {
				name = child.Content(src)
				break
			}
		}
		if name == "" {
			continue
		}

		qualName := name
		if modulePrefix != "" {
			qualName = modulePrefix + "." + name
		}

		id := filePath + "::" + qualName
		if seen[id] {
			continue
		}
		seen[id] = true

		startLine := int(node.StartPoint().Row) + 1
		endLine := int(node.EndPoint().Row) + 1

		result.Nodes = append(result.Nodes, &graph.Node{
			ID: id, Kind: graph.KindType, Name: name,
			FilePath: filePath, StartLine: startLine, EndLine: endLine,
			Language: "ocaml", Meta: map[string]any{"type_flavor": "module"},
		})
		result.Edges = append(result.Edges, &graph.Edge{
			From: fileNode.ID, To: id, Kind: graph.EdgeDefines,
			FilePath: filePath, Line: startLine,
		})

		// Recurse into module body (structure node) for nested definitions.
		for j, _nc := 0, int(binding.NamedChildCount()); j < _nc; j++ {
			child := binding.NamedChild(j)
			if child != nil && (child.Type() == "structure" || child.Type() == "struct_expression") {
				e.walkExtract(child, src, filePath, fileNode, result, seen, qualName)
			}
		}
	}
}

// extractModuleTypeDef handles `module type Name = sig ... end`
func (e *OCamlExtractor) extractModuleTypeDef(
	node *sitter.Node, src []byte, filePath string, fileNode *graph.Node,
	result *parser.ExtractionResult, seen map[string]bool, modulePrefix string,
) {
	name := ""
	for i, _nc := 0, int(node.NamedChildCount()); i < _nc; i++ {
		child := node.NamedChild(i)
		if child != nil && child.Type() == "module_type_name" {
			name = child.Content(src)
			break
		}
	}
	if name == "" {
		return
	}

	qualName := name
	if modulePrefix != "" {
		qualName = modulePrefix + "." + name
	}

	id := filePath + "::" + qualName
	if seen[id] {
		return
	}
	seen[id] = true

	startLine := int(node.StartPoint().Row) + 1

	result.Nodes = append(result.Nodes, &graph.Node{
		ID: id, Kind: graph.KindInterface, Name: name,
		FilePath: filePath, StartLine: startLine, EndLine: int(node.EndPoint().Row) + 1,
		Language: "ocaml", Meta: map[string]any{"type_flavor": "signature"},
	})
	result.Edges = append(result.Edges, &graph.Edge{
		From: fileNode.ID, To: id, Kind: graph.EdgeDefines,
		FilePath: filePath, Line: startLine,
	})
}

// extractOpen handles `open ModuleName`
func (e *OCamlExtractor) extractOpen(
	node *sitter.Node, src []byte, filePath string, fileNode *graph.Node,
	result *parser.ExtractionResult,
) {
	for i, _nc := 0, int(node.NamedChildCount()); i < _nc; i++ {
		child := node.NamedChild(i)
		if child != nil && (child.Type() == "module_path" || child.Type() == "module_name" || child.Type() == "extended_module_path") {
			moduleName := child.Content(src)
			result.Edges = append(result.Edges, &graph.Edge{
				From: fileNode.ID, To: "unresolved::import::" + moduleName,
				Kind: graph.EdgeImports, FilePath: filePath, Line: int(node.StartPoint().Row) + 1,
			})
			return
		}
	}
}

// extractClassDef handles `class name = object ... end`
func (e *OCamlExtractor) extractClassDef(
	node *sitter.Node, src []byte, filePath string, fileNode *graph.Node,
	result *parser.ExtractionResult, seen map[string]bool, modulePrefix string,
) {
	for i, _nc := 0, int(node.NamedChildCount()); i < _nc; i++ {
		child := node.NamedChild(i)
		if child == nil || child.Type() != "class_binding" {
			continue
		}

		name := ""
		for j, _nc := 0, int(child.NamedChildCount()); j < _nc; j++ {
			part := child.NamedChild(j)
			if part != nil && (part.Type() == "class_name" || part.Type() == "value_name") {
				name = part.Content(src)
				break
			}
		}

		if name == "" {
			continue
		}

		qualName := name
		if modulePrefix != "" {
			qualName = modulePrefix + "." + name
		}

		id := filePath + "::" + qualName
		if seen[id] {
			continue
		}
		seen[id] = true

		startLine := int(child.StartPoint().Row) + 1

		result.Nodes = append(result.Nodes, &graph.Node{
			ID: id, Kind: graph.KindType, Name: name,
			FilePath: filePath, StartLine: startLine, EndLine: int(child.EndPoint().Row) + 1,
			Language: "ocaml", Meta: map[string]any{"type_flavor": "class"},
		})
		result.Edges = append(result.Edges, &graph.Edge{
			From: fileNode.ID, To: id, Kind: graph.EdgeDefines,
			FilePath: filePath, Line: startLine,
		})

		// Extract methods from class body.
		e.extractMethods(child, src, filePath, fileNode, result, seen, qualName)
	}
}

// extractMethods extracts method definitions from a class body.
func (e *OCamlExtractor) extractMethods(
	classNode *sitter.Node, src []byte, filePath string, fileNode *graph.Node,
	result *parser.ExtractionResult, seen map[string]bool, className string,
) {
	walkNodes(classNode, func(node *sitter.Node) {
		if node.Type() != "method_definition" {
			return
		}
		name := ""
		for i, _nc := 0, int(node.NamedChildCount()); i < _nc; i++ {
			child := node.NamedChild(i)
			if child != nil && child.Type() == "method_name" {
				name = child.Content(src)
				break
			}
		}
		if name == "" {
			return
		}

		id := filePath + "::" + className + "." + name
		if seen[id] {
			return
		}
		seen[id] = true

		startLine := int(node.StartPoint().Row) + 1

		result.Nodes = append(result.Nodes, &graph.Node{
			ID: id, Kind: graph.KindMethod, Name: name,
			FilePath: filePath, StartLine: startLine, EndLine: int(node.EndPoint().Row) + 1,
			Language: "ocaml",
			Meta:     map[string]any{"receiver": className},
		})
		result.Edges = append(result.Edges, &graph.Edge{
			From: fileNode.ID, To: id, Kind: graph.EdgeDefines,
			FilePath: filePath, Line: startLine,
		})
		result.Edges = append(result.Edges, &graph.Edge{
			From: id, To: filePath + "::" + className, Kind: graph.EdgeMemberOf,
			FilePath: filePath, Line: startLine,
		})
	})
}

// extractExternal handles `external name : type = "c_function"`
func (e *OCamlExtractor) extractExternal(
	node *sitter.Node, src []byte, filePath string, fileNode *graph.Node,
	result *parser.ExtractionResult, seen map[string]bool, modulePrefix string,
) {
	name := ""
	for i, _nc := 0, int(node.NamedChildCount()); i < _nc; i++ {
		child := node.NamedChild(i)
		if child != nil && child.Type() == "value_name" {
			name = child.Content(src)
			break
		}
	}
	if name == "" {
		return
	}

	qualName := name
	if modulePrefix != "" {
		qualName = modulePrefix + "." + name
	}

	id := filePath + "::" + qualName
	if seen[id] {
		return
	}
	seen[id] = true

	startLine := int(node.StartPoint().Row) + 1

	result.Nodes = append(result.Nodes, &graph.Node{
		ID: id, Kind: graph.KindFunction, Name: name,
		FilePath: filePath, StartLine: startLine, EndLine: int(node.EndPoint().Row) + 1,
		Language: "ocaml",
		Meta:     map[string]any{"signature": fmt.Sprintf("external %s : ...", name)},
	})
	result.Edges = append(result.Edges, &graph.Edge{
		From: fileNode.ID, To: id, Kind: graph.EdgeDefines,
		FilePath: filePath, Line: startLine,
	})
}

// extractValueSpec handles `val name : type` in .mli files
func (e *OCamlExtractor) extractValueSpec(
	node *sitter.Node, src []byte, filePath string, fileNode *graph.Node,
	result *parser.ExtractionResult, seen map[string]bool, modulePrefix string,
) {
	name := ""
	for i, _nc := 0, int(node.NamedChildCount()); i < _nc; i++ {
		child := node.NamedChild(i)
		if child != nil && child.Type() == "value_name" {
			name = child.Content(src)
			break
		}
	}
	if name == "" {
		return
	}

	qualName := name
	if modulePrefix != "" {
		qualName = modulePrefix + "." + name
	}

	id := filePath + "::" + qualName
	if seen[id] {
		return
	}
	seen[id] = true

	startLine := int(node.StartPoint().Row) + 1

	result.Nodes = append(result.Nodes, &graph.Node{
		ID: id, Kind: graph.KindFunction, Name: name,
		FilePath: filePath, StartLine: startLine, EndLine: int(node.EndPoint().Row) + 1,
		Language: "ocaml",
		Meta:     map[string]any{"signature": fmt.Sprintf("val %s : ...", name)},
	})
	result.Edges = append(result.Edges, &graph.Edge{
		From: fileNode.ID, To: id, Kind: graph.EdgeDefines,
		FilePath: filePath, Line: startLine,
	})
}

// extractCalls walks the AST for function application nodes.
func (e *OCamlExtractor) extractCalls(
	root *sitter.Node, src []byte, filePath string,
	result *parser.ExtractionResult, funcRanges []funcRange,
) {
	walkNodes(root, func(node *sitter.Node) {
		if node.Type() != "application" {
			return
		}

		// The first child of an application is the function being called.
		if node.NamedChildCount() == 0 {
			return
		}
		funcNode := node.NamedChild(0)
		if funcNode == nil {
			return
		}

		var callName string
		switch funcNode.Type() {
		case "value_path", "value_name":
			callName = funcNode.Content(src)
		case "field_get_expression":
			// Module.function or record.field
			callName = funcNode.Content(src)
		default:
			return
		}

		if callName == "" {
			return
		}

		line := int(node.StartPoint().Row) + 1
		callerID := findEnclosingFunc(funcRanges, line)
		if callerID == "" {
			return
		}

		// Check if it's a qualified call (Module.function).
		if strings.Contains(callName, ".") {
			parts := strings.SplitN(callName, ".", 2)
			if len(parts) == 2 {
				result.Edges = append(result.Edges, &graph.Edge{
					From: callerID, To: "unresolved::*." + parts[1],
					Kind: graph.EdgeCalls, FilePath: filePath, Line: line,
				})
				return
			}
		}

		result.Edges = append(result.Edges, &graph.Edge{
			From: callerID, To: "unresolved::" + callName,
			Kind: graph.EdgeCalls, FilePath: filePath, Line: line,
		})
	})
}