Language.java

package io.github.treesitter.jtreesitter;

import static io.github.treesitter.jtreesitter.internal.TreeSitter.*;

import io.github.treesitter.jtreesitter.internal.TreeSitter;
import java.lang.foreign.*;
import java.util.Objects;
import org.jspecify.annotations.NullMarked;
import org.jspecify.annotations.Nullable;

/** A class that defines how to parse a particular language. */
@NullMarked
public final class Language {
    /**
     * The latest ABI version that is supported by the current version of the library.
     *
     * @apiNote The Tree-sitter library is generally backwards-compatible with
     * languages generated using older CLI versions, but is not forwards-compatible.
     */
    public static final @Unsigned int LANGUAGE_VERSION = TREE_SITTER_LANGUAGE_VERSION();

    /** The earliest ABI version that is supported by the current version of the library. */
    public static final @Unsigned int MIN_COMPATIBLE_LANGUAGE_VERSION = TREE_SITTER_MIN_COMPATIBLE_LANGUAGE_VERSION();

    private static final ValueLayout VOID_PTR =
            ValueLayout.ADDRESS.withTargetLayout(MemoryLayout.sequenceLayout(Long.MAX_VALUE, ValueLayout.JAVA_BYTE));

    private static final FunctionDescriptor FUNC_DESC = FunctionDescriptor.of(VOID_PTR);

    private static final Linker LINKER = Linker.nativeLinker();

    private final MemorySegment self;

    private final @Unsigned int version;

    /**
     * Creates a new instance from the given language pointer.
     *
     * @implNote It is up to the caller to ensure that the pointer is valid.
     *
     * @throws IllegalArgumentException If the language version is incompatible.
     */
    public Language(MemorySegment self) throws IllegalArgumentException {
        this.self = Objects.requireNonNull(self);
        version = ts_language_version(this.self);
        if (version < MIN_COMPATIBLE_LANGUAGE_VERSION || version > LANGUAGE_VERSION) {
            throw new IllegalArgumentException(String.format(
                    "Incompatible language version %d. Must be between %d and %d.",
                    version, MIN_COMPATIBLE_LANGUAGE_VERSION, LANGUAGE_VERSION));
        }
    }

    private static UnsatisfiedLinkError unresolved(String name) {
        return new UnsatisfiedLinkError("Unresolved symbol: %s".formatted(name));
    }

    private static MemorySegment loadLanguagePointer(SymbolLookup symbols, Arena libraryArena, String language) {
        var address = symbols.find(language).orElseThrow(() -> unresolved(language));
        try {
            var function = LINKER.downcallHandle(address, FUNC_DESC);
            var languagePointer = (MemorySegment) function.invokeExact();
            // The results of Linker downcalls always use the global scope, but the language pointer actually points
            // to data in the loaded library. Therefore change the scope of the pointer to be the same as the library.
            // So if the library is unloaded while the language pointer is still in use, the language pointer becomes
            // invalid and an exception occurs (instead of a JVM crash).
            // Ideally this would not require the Arena (which is not always available), but instead just apply the
            // scope of `address` to the `languagePointer`, but that is currently not possible, see https://bugs.openjdk.org/browse/JDK-8340641
            return languagePointer.reinterpret(libraryArena, TreeSitter::ts_language_delete);
        } catch (Throwable e) {
            throw new RuntimeException("Failed to load %s".formatted(language), e);
        }
    }

    /**
     * Load a language by looking for its function in the given symbols.
     *
     * <h4 id="load-example">Example</h4>
     *
     * <p>{@snippet lang="java" :
     * String library = System.mapLibraryName("tree-sitter-java");
     * Arena arena = Arena.global();
     * SymbolLookup symbols = SymbolLookup.libraryLookup(library, arena);
     * Language language = Language.load(symbols, arena, "tree_sitter_java");
     * }
     *
     * @param libraryArena The arena which was used to load the {@code symbols} library.
     * @throws UnsatisfiedLinkError If the language function could not be found in the symbols.
     * @throws RuntimeException If the language could not be loaded.
     */
    // TODO: deprecate when the bindings are generated by the CLI
    public static Language load(SymbolLookup symbols, Arena libraryArena, String language) throws RuntimeException {
        return new Language(loadLanguagePointer(symbols, libraryArena, language));
    }

    /**
     * Load a language by looking for its function in the given symbols.
     *
     * <p>If {@code symbols} was obtained from {@link SymbolLookup#libraryLookup} and you have access to the
     * {@code Arena} which was used as argument to {@code libraryLookup(...)}, prefer
     * {@link #load(SymbolLookup, Arena, String)} since it is safer.
     *
     * <h4 id="load-example">Example</h4>
     *
     * <p>{@snippet lang="java" :
     * SymbolLookup symbols = ...;
     * Language language = Language.load(symbols, "tree_sitter_java");
     * }
     *
     * @throws UnsatisfiedLinkError If the language function could not be found in the symbols.
     * @throws RuntimeException If the language could not be loaded.
     */
    // TODO: deprecate when the bindings are generated by the CLI
    public static Language load(SymbolLookup symbols, String language) throws RuntimeException {
        // TODO: This is not actually safe, `libraryArena` should be the arena which was used to create `symbols`,
        //   but we don't have access to it here; see comments in `loadLanguagePointer(...)` for details
        //   In the worst case the JVM crashes if the user unloads the parser library while the language is still in use
        var libraryArena = LIBRARY_ARENA;
        return new Language(loadLanguagePointer(symbols, libraryArena, language));
    }

    MemorySegment segment() {
        return self;
    }

    /**
     * Get the ABI version number for this language.
     *
     * <p>When a language is generated by the Tree-sitter CLI, it is assigned
     * an ABI version number that corresponds to the current CLI version.
     */
    public @Unsigned int getVersion() {
        return version;
    }

    /** Get the number of distinct node types in this language. */
    public @Unsigned int getSymbolCount() {
        return ts_language_symbol_count(self);
    }

    /** Get the number of valid states in this language */
    public @Unsigned int getStateCount() {
        return ts_language_state_count(self);
    }

    /** Get the number of distinct field names in this language */
    public @Unsigned int getFieldCount() {
        return ts_language_field_count(self);
    }

    /** Get the node type for the given numerical ID. */
    public @Nullable String getSymbolName(@Unsigned short symbol) {
        var name = ts_language_symbol_name(self, symbol);
        return name.equals(MemorySegment.NULL) ? null : name.getString(0);
    }

    /** Get the numerical ID for the given node type, or {@code 0} if not found. */
    public @Unsigned short getSymbolForName(String name, boolean isNamed) {
        try (var arena = Arena.ofConfined()) {
            var segment = arena.allocateFrom(name);
            return ts_language_symbol_for_name(self, segment, name.length(), isNamed);
        }
    }

    /**
     * Check if the node for the given numerical ID is named.
     *
     * @see Node#isNamed
     */
    public boolean isNamed(@Unsigned short symbol) {
        return ts_language_symbol_type(self, symbol) == TSSymbolTypeRegular();
    }

    /** Check if the node for the given numerical ID is visible. */
    public boolean isVisible(@Unsigned short symbol) {
        return ts_language_symbol_type(self, symbol) <= TSSymbolTypeAnonymous();
    }

    /** Get the field name for the given numerical id. */
    public @Nullable String getFieldNameForId(@Unsigned short id) {
        var name = ts_language_field_name_for_id(self, id);
        return name.equals(MemorySegment.NULL) ? null : name.getString(0);
    }

    /** Get the numerical ID for the given field name. */
    public @Unsigned short getFieldIdForName(String name) {
        try (var arena = Arena.ofConfined()) {
            var segment = arena.allocateFrom(name);
            return ts_language_field_id_for_name(self, segment, name.length());
        }
    }

    /**
     * Get the next parse state.
     *
     * <p>{@snippet lang="java" :
     * short state = language.nextState(node.getParseState(), node.getGrammarSymbol());
     * }
     *
     * <p>Combine this with {@link #lookaheadIterator lookaheadIterator(state)}
     * to generate completion suggestions or valid symbols in {@index ERROR} nodes.
     */
    public @Unsigned short nextState(@Unsigned short state, @Unsigned short symbol) {
        return ts_language_next_state(self, state, symbol);
    }

    /**
     * Create a new lookahead iterator for the given parse state.
     *
     * @throws IllegalArgumentException If the state is invalid for this language.
     */
    public LookaheadIterator lookaheadIterator(@Unsigned short state) throws IllegalArgumentException {
        return new LookaheadIterator(self, state);
    }

    /**
     * Create a new query from a string containing one or more S-expression patterns.
     *
     * @throws QueryError If an error occurred while creating the query.
     */
    public Query query(String source) throws QueryError {
        return new Query(this, source);
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (!(o instanceof Language other)) return false;
        return self.equals(other.self);
    }

    @Override
    public int hashCode() {
        return Long.hashCode(self.address());
    }

    @Override
    public String toString() {
        return "Language{id=0x%x, version=%d}".formatted(self.address(), version);
    }
}