#include "ir.h"

#include "mem.h"

#include <stdlib.h>
#include <stdint.h>

tobi_ir *tobi_ir_new(tobi_ir_kind kind, size_t off, size_t len) {
    tobi_ir *n = tobi_xcalloc(1, sizeof(*n));
    n->kind = kind;
    n->off = off;
    n->len = len;
    n->parent_index = SIZE_MAX;
    return n;
}

void tobi_ir_free(tobi_ir *node) {
    if (!node) {
        return;
    }
    for (size_t i = 0; i < node->child_len; i++) {
        tobi_ir_free(node->child[i]);
    }
    free(node->child);
    free(node->name);
    free(node->path);
    free(node->scope);
    free(node->str);
    free(node->target);
    free(node->source);
    free(node);
}

void tobi_ir_add(tobi_ir *parent, tobi_ir *child) {
    if (!parent || !child) {
        return;
    }
    if (parent->child_len == parent->child_cap) {
        parent->child_cap = tobi_xgrow_cap(parent->child_cap, parent->child_len + 1u, 4u);
        parent->child = tobi_xrealloc(parent->child,
                                      tobi_xmul_size(parent->child_cap, sizeof(parent->child[0])));
    }
    if (parent->has_source && !child->has_source) {
        tobi_ir_set_source_recursive(child, parent->source, parent->input_index);
    }
    child->parent_id = parent->id;
    child->parent_index = parent->child_len;
    parent->child[parent->child_len++] = child;
}

void tobi_ir_assign_ids(tobi_ir *root) {
    if (!root) {
        return;
    }
    typedef struct frame {
        tobi_ir *node;
        size_t parent_id;
        size_t parent_index;
    } frame;
    size_t cap = 64;
    size_t len = 0;
    frame *stack = tobi_xcalloc(cap, sizeof(stack[0]));
    stack[len++] = (frame){root, 0, SIZE_MAX};
    size_t next_id = 1;
    while (len > 0) {
        frame cur = stack[--len];
        cur.node->id = next_id++;
        cur.node->parent_id = cur.parent_id;
        cur.node->parent_index = cur.parent_index;
        for (size_t i = cur.node->child_len; i > 0; i--) {
            if (len == cap) {
                cap = tobi_xgrow_cap(cap, len + 1u, 64u);
                stack = tobi_xrealloc(stack, tobi_xmul_size(cap, sizeof(stack[0])));
            }
            stack[len++] = (frame){cur.node->child[i - 1u], cur.node->id, i - 1u};
        }
    }
    free(stack);
}

void tobi_ir_set_name(tobi_ir *node, const char *s) {
    free(node->name);
    node->name = tobi_xstrdup(s ? s : "");
}

void tobi_ir_set_path(tobi_ir *node, const char *s) {
    free(node->path);
    node->path = tobi_xstrdup(s ? s : "");
}

void tobi_ir_set_scope(tobi_ir *node, const char *s) {
    free(node->scope);
    node->scope = tobi_xstrdup(s ? s : "");
}

void tobi_ir_set_str(tobi_ir *node, const char *s) {
    free(node->str);
    node->str = tobi_xstrdup(s ? s : "");
}

void tobi_ir_set_target(tobi_ir *node, const char *s) {
    free(node->target);
    node->target = tobi_xstrdup(s ? s : "");
}

void tobi_ir_set_source(tobi_ir *node, const char *s, size_t input_index) {
    if (!node || !s) {
        return;
    }
    free(node->source);
    node->source = tobi_xstrdup(s);
    node->input_index = input_index;
    node->has_source = 1;
}

void tobi_ir_set_source_recursive(tobi_ir *node, const char *s, size_t input_index) {
    if (!node || !s) {
        return;
    }
    tobi_ir_set_source(node, s, input_index);
    for (size_t i = 0; i < node->child_len; i++) {
        tobi_ir_set_source_recursive(node->child[i], s, input_index);
    }
}

const char *tobi_ir_kind_name(tobi_ir_kind kind) {
    switch (kind) {
    case TOBI_IR_ROOT: return "root";
    case TOBI_IR_BLOCK: return "block";
    case TOBI_IR_SCOPE: return "scope";
    case TOBI_IR_DEVICE: return "device";
    case TOBI_IR_METHOD: return "method";
    case TOBI_IR_PROCESSOR: return "processor";
    case TOBI_IR_NAME: return "name";
    case TOBI_IR_ALIAS: return "alias";
    case TOBI_IR_OPREGION: return "opregion";
    case TOBI_IR_FIELD: return "field";
    case TOBI_IR_FIELD_ELEM: return "field_elem";
    case TOBI_IR_MUTEX: return "mutex";
    case TOBI_IR_EVENT: return "event";
    case TOBI_IR_RESOURCE: return "resource";
    case TOBI_IR_BUFFER: return "buffer";
    case TOBI_IR_PACKAGE: return "package";
    case TOBI_IR_INTEGER: return "integer";
    case TOBI_IR_STRING: return "string";
    case TOBI_IR_REF: return "ref";
    case TOBI_IR_CALL: return "call";
    case TOBI_IR_EXPR: return "expr";
    case TOBI_IR_STORE: return "store";
    case TOBI_IR_RETURN: return "return";
    case TOBI_IR_IF: return "if";
    case TOBI_IR_WHILE: return "while";
    case TOBI_IR_BREAK: return "break";
    case TOBI_IR_CONTINUE: return "continue";
    case TOBI_IR_UNKNOWN: return "unknown";
    case TOBI_IR_DIAG: return "diag";
    }
    return "invalid";
}

tobi_ir *tobi_ir_find_child(tobi_ir *node, tobi_ir_kind kind) {
    if (!node) {
        return NULL;
    }
    for (size_t i = 0; i < node->child_len; i++) {
        if (node->child[i]->kind == kind) {
            return node->child[i];
        }
    }
    return NULL;
}

size_t tobi_ir_count_kind(const tobi_ir *node, tobi_ir_kind kind) {
    size_t n = 0;
    if (!node) {
        return 0;
    }
    for (size_t i = 0; i < node->child_len; i++) {
        if (node->child[i]->kind == kind) {
            n++;
        }
    }
    return n;
}