nub-lang-archive/compiler/NubLang.LSP/AstExtensions.cs

using NubLang.Ast;

namespace NubLang.LSP;

public static class AstExtensions
{
    public static bool ContainsPosition(this Node node, int line, int character)
    {
        if (node.Tokens.Count == 0)
        {
            return false;
        }

        var start = node.Tokens.First().Span.Start;
        var end = node.Tokens.Last().Span.End;

        var startLine = start.Line - 1;
        var startChar = start.Column - 1;
        var endLine = end.Line - 1;
        var endChar = end.Column - 1;

        if (line < startLine || line > endLine) return false;

        if (line > startLine && line < endLine) return true;

        if (startLine == endLine)
        {
            return character >= startChar && character <= endChar;
        }

        if (line == startLine)
        {
            return character >= startChar;
        }

        if (line == endLine)
        {
            return character <= endChar;
        }

        return false;
    }

    public static FuncNode? FunctionAtPosition(this CompilationUnit compilationUnit, int line, int character)
    {
        return compilationUnit
            .Functions
            .FirstOrDefault(x => x.ContainsPosition(line, character));
    }

    public static Node? DeepestNodeAtPosition(this CompilationUnit compilationUnit, int line, int character)
    {
        return compilationUnit.Functions
            .SelectMany(x => x.EnumerateDescendantsAndSelf())
            .Where(n => n.ContainsPosition(line, character))
            .OrderBy(n => n.Tokens.First().Span.Start.Line)
            .ThenBy(n => n.Tokens.First().Span.Start.Column)
            .LastOrDefault();
    }

    public static Node? DeepestNodeAtPosition(this Node node, int line, int character)
    {
        return node.EnumerateDescendantsAndSelf()
            .Where(n => n.ContainsPosition(line, character))
            .OrderBy(n => n.Tokens.First().Span.Start.Line)
            .ThenBy(n => n.Tokens.First().Span.Start.Column)
            .LastOrDefault();
    }

    public static IEnumerable<Node> EnumerateDescendantsAndSelf(this Node node)
    {
        yield return node;

        switch (node)
        {
            case FuncNode func:
            {
                foreach (var n in func.Prototype.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                if (func.Body != null)
                {
                    foreach (var n in func.Body.EnumerateDescendantsAndSelf())
                    {
                        yield return n;
                    }
                }

                break;
            }
            case FuncPrototypeNode proto:
            {
                foreach (var n in proto.Parameters.SelectMany(param => param.EnumerateDescendantsAndSelf()))
                {
                    yield return n;
                }

                break;
            }
            case BlockNode block:
            {
                foreach (var n in block.Statements.SelectMany(stmt => stmt.EnumerateDescendantsAndSelf()))
                {
                    yield return n;
                }

                break;
            }
            case StatementFuncCallNode stmtCall:
            {
                foreach (var n in stmtCall.FuncCall.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case ReturnNode { Value: not null } ret:
            {
                foreach (var n in ret.Value.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case AssignmentNode assign:
            {
                foreach (var n in assign.Target.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                foreach (var n in assign.Value.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case IfNode ifNode:
            {
                foreach (var n in ifNode.Condition.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                foreach (var n in ifNode.Body.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                if (ifNode.Else.HasValue)
                {
                    if (ifNode.Else.Value.IsCase1(out var elseIfNode))
                    {
                        foreach (var n in elseIfNode.EnumerateDescendantsAndSelf())
                        {
                            yield return n;
                        }
                    }
                    else if (ifNode.Else.Value.IsCase2(out var elseNode))
                    {
                        foreach (var n in elseNode.EnumerateDescendantsAndSelf())
                        {
                            yield return n;
                        }
                    }
                }

                break;
            }
            case VariableDeclarationNode decl:
            {
                if (decl.Assignment != null)
                {
                    foreach (var n in decl.Assignment.EnumerateDescendantsAndSelf())
                    {
                        yield return n;
                    }
                }

                break;
            }
            case WhileNode whileNode:
            {
                foreach (var n in whileNode.Condition.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                foreach (var n in whileNode.Body.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case ForSliceNode forSlice:
            {
                foreach (var n in forSlice.Target.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                foreach (var n in forSlice.Body.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case ForConstArrayNode forConst:
            {
                foreach (var n in forConst.Target.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                foreach (var n in forConst.Body.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case DeferNode defer:
            {
                foreach (var n in defer.Statement.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case BinaryExpressionNode bin:
            {
                foreach (var n in bin.Left.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                foreach (var n in bin.Right.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case UnaryExpressionNode unary:
            {
                foreach (var n in unary.Operand.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case FuncCallNode call:
            {
                foreach (var n in call.Expression.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                foreach (var n in call.Parameters.SelectMany(param => param.EnumerateDescendantsAndSelf()))
                {
                    yield return n;
                }

                break;
            }
            case ArrayInitializerNode arrInit:
            {
                foreach (var n in arrInit.Values.SelectMany(val => val.EnumerateDescendantsAndSelf()))
                {
                    yield return n;
                }

                break;
            }
            case ConstArrayInitializerNode constArrInit:
            {
                foreach (var n in constArrInit.Values.SelectMany(val => val.EnumerateDescendantsAndSelf()))
                {
                    yield return n;
                }

                break;
            }
            case ArrayIndexAccessNode arrIndex:
            {
                foreach (var n in arrIndex.Target.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                foreach (var n in arrIndex.Index.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case ConstArrayIndexAccessNode constArrIndex:
            {
                foreach (var n in constArrIndex.Target.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                foreach (var n in constArrIndex.Index.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case SliceIndexAccessNode sliceIndex:
            {
                foreach (var n in sliceIndex.Target.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                foreach (var n in sliceIndex.Index.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case AddressOfNode addr:
            {
                foreach (var n in addr.LValue.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case StructFieldAccessNode field:
            {
                foreach (var n in field.Target.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case StructInitializerNode structInit:
            {
                foreach (var n in structInit.Initializers.SelectMany(kv => kv.Value.EnumerateDescendantsAndSelf()))
                {
                    yield return n;
                }

                break;
            }
            case DereferenceNode deref:
            {
                foreach (var n in deref.Target.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case ConvertIntNode convInt:
            {
                foreach (var n in convInt.Value.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case ConvertFloatNode convFloat:
            {
                foreach (var n in convFloat.Value.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case ConvertCStringToStringNode convStr:
            {
                foreach (var n in convStr.Value.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case FloatToIntBuiltinNode ftoi:
            {
                foreach (var n in ftoi.Value.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
            case ConstArrayToSliceNode constSlice:
            {
                foreach (var n in constSlice.Array.EnumerateDescendantsAndSelf())
                {
                    yield return n;
                }

                break;
            }
        }
    }
}