wraith-lang/src/identcheck.c

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#include "ast.h"
#include "identcheck.h"

IdNode* MakeIdNode(NodeType type, char *name, IdNode *parent) {
    IdNode *node = (IdNode*)malloc(sizeof(IdNode));
    node->type = type;
    node->name = strdup(name);
    node->parent = parent;
    node->childCount = 0;
    node->childCapacity = 0;
    node->children = NULL;
    node->typeTag = NULL;
    return node;
}

void AddChildToNode(IdNode *node, IdNode *child) {
    if (child == NULL) return;

    if (node->children == NULL) {
        node->childCapacity = 2;
        node->children = (IdNode**) malloc(sizeof(IdNode*) * node->childCapacity);
    } else if (node->childCount == node->childCapacity) {
        node->childCapacity *= 2;
        node->children = (IdNode**) realloc(node->children, sizeof(IdNode*) * node->childCapacity);
    }

    node->children[node->childCount] = child;
    node->childCount += 1;
}

IdNode* MakeIdTree(Node *astNode, IdNode *parent) {
    uint32_t i;
    IdNode *mainNode;
    switch (astNode->syntaxKind) {
        case AccessExpression:
            AddChildToNode(parent, MakeIdTree(astNode->accessExpression.accessee, parent));
            AddChildToNode(parent, MakeIdTree(astNode->accessExpression.accessor, parent));
            return NULL;

        case AllocExpression:
            AddChildToNode(parent, MakeIdTree(astNode->allocExpression.type, parent));
            return NULL;

        case Assignment: {
            if (astNode->assignmentStatement.left->syntaxKind == Declaration) {
                return MakeIdTree(astNode->assignmentStatement.left, parent);
            } else {
                AddChildToNode(parent, MakeIdTree(astNode->assignmentStatement.left, parent));
                AddChildToNode(parent, MakeIdTree(astNode->assignmentStatement.right, parent));
                return NULL;
            }
        }

        case BinaryExpression:
            AddChildToNode(parent, MakeIdTree(astNode->binaryExpression.left, parent));
            AddChildToNode(parent, MakeIdTree(astNode->binaryExpression.right, parent));
            return NULL;

        case Declaration: {
            Node *idNode = astNode->declaration.identifier;
            mainNode = MakeIdNode(Variable, idNode->identifier.name, parent);
            mainNode->typeTag = MakeTypeTag(astNode);
            idNode->typeTag = mainNode->typeTag;
            break;
        }

        case DeclarationSequence: {
            mainNode = MakeIdNode(UnorderedScope, "", parent);
            for (i = 0; i < astNode->declarationSequence.count; i++) {
                AddChildToNode(
                    mainNode, MakeIdTree(astNode->declarationSequence.sequence[i], mainNode));
            }
            break;
        }

        case ForLoop: {
            Node *loopDecl = astNode->forLoop.declaration;
            Node *loopBody = astNode->forLoop.statementSequence;
            mainNode = MakeIdNode(OrderedScope, "for-loop", parent);
            AddChildToNode(mainNode, MakeIdTree(loopDecl, mainNode));
            AddChildToNode(mainNode, MakeIdTree(loopBody, mainNode));
            break;
        }

        case FunctionArgumentSequence:
            for (i = 0; i < astNode->functionArgumentSequence.count; i++) {
                AddChildToNode(
                    parent, MakeIdTree(astNode->functionArgumentSequence.sequence[i], parent));
            }
            return NULL;

        case FunctionCallExpression:
            AddChildToNode(parent, MakeIdTree(astNode->functionCallExpression.identifier, parent));
            AddChildToNode(
                parent, MakeIdTree(astNode->functionCallExpression.argumentSequence, parent));
            return NULL;

        case FunctionDeclaration: {
            Node *sigNode = astNode->functionDeclaration.functionSignature;
            Node *idNode = sigNode->functionSignature.identifier;
            char *funcName = idNode->identifier.name;
            mainNode = MakeIdNode(Function, funcName, parent);
            mainNode->typeTag = MakeTypeTag(astNode);
            idNode->typeTag = mainNode->typeTag;
            MakeIdTree(sigNode->functionSignature.arguments, mainNode);
            MakeIdTree(astNode->functionDeclaration.functionBody, mainNode);
            break;
        }

        case FunctionSignatureArguments: {
            for (i = 0; i < astNode->functionSignatureArguments.count; i++) {
                Node *argNode = astNode->functionSignatureArguments.sequence[i];
                AddChildToNode(parent, MakeIdTree(argNode, parent));
            }
            return NULL;
        }

        case Identifier: {
            char *name = astNode->identifier.name;
            mainNode = MakeIdNode(Placeholder, name, parent);
            IdNode *lookupNode = LookupId(mainNode, NULL, name);
            if (lookupNode == NULL) {
                fprintf(stderr, "wraith: Could not find IdNode for id %s\n", name);
                TypeTag *tag = (TypeTag*)malloc(sizeof(TypeTag));
                tag->type = Unknown;
                astNode->typeTag = tag;
            } else {
                astNode->typeTag = lookupNode->typeTag;
            }
            break;
        }

        case IfStatement: {
            Node *clause = astNode->ifStatement.expression;
            Node *stmtSeq = astNode->ifStatement.statementSequence;
            mainNode = MakeIdNode(OrderedScope, "if", parent);
            MakeIdTree(clause, mainNode);
            MakeIdTree(stmtSeq, mainNode);
            break;
        }

        case IfElseStatement: {
            Node *ifNode = astNode->ifElseStatement.ifStatement;
            Node *elseStmts = astNode->ifElseStatement.elseStatement;
            mainNode = MakeIdNode(OrderedScope, "if-else", parent);
            IdNode *ifBranch = MakeIdTree(ifNode, mainNode);
            AddChildToNode(mainNode, ifBranch);
            IdNode *elseScope = MakeIdNode(OrderedScope, "else", mainNode);
            MakeIdTree(elseStmts, elseScope);
            AddChildToNode(mainNode, elseScope);
            break;
        }

        case ReferenceTypeNode:
            AddChildToNode(parent, MakeIdTree(astNode->referenceType.type, parent));
            return NULL;

        case Return:
            AddChildToNode(parent, MakeIdTree(astNode->returnStatement.expression, parent));
            return NULL;

        case StatementSequence: {
            for (i = 0; i < astNode->statementSequence.count; i++) {
                Node *argNode = astNode->statementSequence.sequence[i];
                AddChildToNode(parent, MakeIdTree(argNode, parent));
            }
            return NULL;
        }

        case StructDeclaration: {
            Node *idNode = astNode->structDeclaration.identifier;
            Node *declsNode = astNode->structDeclaration.declarationSequence;
            mainNode = MakeIdNode(Struct, idNode->identifier.name, parent);
            mainNode->typeTag = MakeTypeTag(astNode);
            for (i = 0; i < declsNode->declarationSequence.count; i++) {
                Node *decl = declsNode->declarationSequence.sequence[i];
                AddChildToNode(mainNode, MakeIdTree(decl, mainNode));
            }
            break;
        }

        case Type:
            AddChildToNode(parent, MakeIdTree(astNode->type.typeNode, parent));
            return NULL;

        case UnaryExpression:
            AddChildToNode(parent, MakeIdTree(astNode->unaryExpression.child, parent));
            return NULL;

        case Comment:
        case CustomTypeNode:
        case FunctionModifiers:
        case FunctionSignature:
        case Number:
        case PrimitiveTypeNode:
        case ReturnVoid:
        case StaticModifier:
        case StringLiteral:
            return NULL;
    }

    astNode->idLink = mainNode;
    return mainNode;
}

void PrintIdNode(IdNode *node) {
    if (node == NULL) {
        fprintf(stderr, "wraith: Attempted to call PrintIdNode with null value.\n");
        return;
    }

    switch(node->type) {
        case Placeholder:
            printf("Placeholder (%s)\n", node->name);
            break;
        case OrderedScope:
            printf("OrderedScope (%s)\n", node->name);
            break;
        case UnorderedScope:
            printf("UnorderedScope (%s)\n", node->name);
            break;
        case Struct:
            printf("%s : %s\n", node->name, TypeTagToString(node->typeTag));
            break;
        case Function:
            printf("%s : Function<%s>\n", node->name, TypeTagToString(node->typeTag));
            break;
        case Variable:
            printf("%s : %s\n", node->name, TypeTagToString(node->typeTag));
            break;
    }
}

void PrintIdTree(IdNode *tree, uint32_t tabCount) {
    if (tree == NULL) {
        fprintf(stderr, "wraith: Attempted to call PrintIdTree on a null value.\n");
        return;
    }

    uint32_t i;
    for (i = 0; i < tabCount; i++) {
        printf("| ");
    }

    PrintIdNode(tree);

    for (i = 0; i < tree->childCount; i++) {
        PrintIdTree(tree->children[i], tabCount + 1);
    }
}

int PrintAncestors(IdNode *node) {
    if (node == NULL) return -1;

    int i;
    int indent = 1;
    indent += PrintAncestors(node->parent);
    for (i = 0; i < indent; i++) {
        printf(" ");
    }
    PrintIdNode(node);
    return indent;
}

IdNode* LookdownId(IdNode *root, NodeType targetType, char *targetName) {
    if (root == NULL) {
        fprintf(stderr, "wraith: Attempted to call LookdownId on a null value.\n");
        return NULL;
    }

    IdNode *result = NULL;
    IdNode **frontier = (IdNode**)malloc(sizeof(IdNode*));
    frontier[0] = root;
    uint32_t frontierCount = 1;

    while (frontierCount > 0) {
        IdNode *current = frontier[0];

        if (current->type == targetType && strcmp(current->name, targetName) == 0) {
            result = current;
            break;
        }

        uint32_t i;
        for(i = 1; i < frontierCount; i++) {
            frontier[i-1] = frontier[i];
        }
        size_t newSize = frontierCount + current->childCount - 1;
        if (frontierCount != newSize) {
            frontier = (IdNode**) realloc(frontier, sizeof(IdNode*) * newSize);
        }
        for (i = 0; i < current->childCount; i++) {
            frontier[frontierCount + i - 1] = current->children[i];
        }
        frontierCount = newSize;
    }

    free(frontier);
    return result;
}

bool ScopeHasOrdering(IdNode *node) {
    switch (node->type) {
        case OrderedScope:
        case Function:
        case Variable: /* this is only technically true */
            return true;
        default:
            return false;
    }
}

IdNode* LookupId(IdNode *node, IdNode *prev, char *target) {
    if (node == NULL) {
        return NULL;
    }

    if (strcmp(node->name, target) == 0 && node->type != Placeholder) {
        return node;
    }

    /* If this is the start of our search, we should not attempt to look at child nodes. Only
     * looking up the scope tree is valid at this point.
     *
     * This has the notable side-effect that this function will return NULL if you attempt to look
     * up a struct's internals starting from the node representing the struct itself. This is
     * because an IdNode corresponds to the location *where an identifier is first declared.* Thus,
     * an identifier has no knowledge of identifiers declared "inside" of it.
     */
    if (prev == NULL) {
        return LookupId(node->parent, node, target);
    }

    /* If the current node forms an ordered scope then we want to prevent ourselves from looking
     * up identifiers declared after the scope we have just come from.
     */
    uint32_t idxLimit;
    if (ScopeHasOrdering(node)) {
        uint32_t i;
        for (i = 0, idxLimit = 0; i < node->childCount; i++, idxLimit++) {
            if (node->children[i] == prev) {
                break;
            }
        }
    } else {
        idxLimit = node->childCount;
    }

    uint32_t i;
    for (i = 0; i < idxLimit; i++) {
        IdNode *child = node->children[i];
        if (child == prev || child->type == Placeholder) {
            /* Do not inspect the node we just came from or placeholders. */
            continue;
        }

        if (strcmp(child->name, target) == 0) {
            return child;
        }

        if (child->type == Struct) {
            uint32_t j;
            for (j = 0; j < child->childCount; j++) {
                IdNode *grandchild = child->children[j];
                if (strcmp(grandchild->name, target) == 0) {
                    return grandchild;
                }
            }
        }
    }

    return LookupId(node->parent, node, target);
}