Basic register allocator in C

regalloc.c

// Copyright (c) 2024 Elric Neumann. All rights reserved. MIT license.
#include <limits.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>

#define MAX_NODES 1000
#define MAX_COLORS 32
#define MAX_NEIGHBORS 100

struct graph_node {
  int id;
  int degree;
  int color;
  bool marked;
  struct graph_node* neighbors[MAX_NEIGHBORS];
  int neighbor_count;
};

struct graph {
  struct graph_node* nodes[MAX_NODES];
  int node_count;
};

struct stack {
  struct graph_node* data[MAX_NODES];
  int top;
};

void stack_push(struct stack* stack, struct graph_node* node) {
  if (stack->top < MAX_NODES - 1) {
    stack->data[++stack->top] = node;
  }
}

struct graph_node* stack_pop(struct stack* stack) {
  if (stack->top >= 0) {
    return stack->data[stack->top--];
  }

  return NULL;
}

struct graph_node* find_or_create_node(struct graph* graph, int id) {
  for (int i = 0; i < graph->node_count; i++) {
    if (graph->nodes[i]->id == id) {
      return graph->nodes[i];
    }
  }

  if (graph->node_count >= MAX_NODES) {
    return NULL;
  }

  struct graph_node* new_node = malloc(sizeof(struct graph_node));

  new_node->id = id;
  new_node->degree = 0;
  new_node->color = -1;
  new_node->marked = false;
  new_node->neighbor_count = 0;

  graph->nodes[graph->node_count++] = new_node;

  return new_node;
}

void add_edge(struct graph* graph, int src, int dest) {
  struct graph_node* src_node = find_or_create_node(graph, src);
  struct graph_node* dest_node = find_or_create_node(graph, dest);

  if (!src_node || !dest_node) {
    return;
  }

  if (src_node->neighbor_count < MAX_NEIGHBORS &&
      dest_node->neighbor_count < MAX_NEIGHBORS) {
    src_node->neighbors[src_node->neighbor_count++] = dest_node;
    dest_node->neighbors[dest_node->neighbor_count++] = src_node;
    src_node->degree++;
    dest_node->degree++;
  }
}

struct graph_node* find_lowest_degree_node(struct graph* graph) {
  struct graph_node* lowest_degree = NULL;
  int min_degree = INT_MAX;

  for (int i = 0; i < graph->node_count; i++) {
    struct graph_node* node = graph->nodes[i];

    if (!node->marked && node->degree < min_degree) {
      lowest_degree = node;
      min_degree = node->degree;
    }
  }

  return lowest_degree;
}

bool can_color(struct graph_node* node, int color) {
  for (int i = 0; i < node->neighbor_count; i++) {
    if (node->neighbors[i]->color == color) {
      return false;
    }
  }

  return true;
}

int color_graph(struct graph* graph, int num_registers) {
  struct stack node_stack = {0};

  while (true) {
    struct graph_node* node = find_lowest_degree_node(graph);
    if (!node) break;

    node->marked = true;
    stack_push(&node_stack, node);
  }

  while (node_stack.top >= 0) {
    struct graph_node* node = stack_pop(&node_stack);

    for (int color = 0; color < num_registers; color++) {
      if (can_color(node, color)) {
        node->color = color;
        break;
      }
    }

    if (node->color == -1) {
      return 0; /* fails to allocate registers */
    }
  }

  return 1; /* successful reg alloc */
}

void free_graph(struct graph* graph) {
  for (int i = 0; i < graph->node_count; i++) {
    free(graph->nodes[i]);
  }
}

int main() {
  struct graph graph = {0};

  add_edge(&graph, 1, 2);
  add_edge(&graph, 1, 3);
  add_edge(&graph, 2, 3);

  int reg_count = 2;
  int output = color_graph(&graph, reg_count);

  printf("reg alloc: %s\n", output ? "pass" : "fail");

  for (int i = 0; i < graph.node_count; i++) {
    printf("node %d: reg %d\n", graph.nodes[i]->id, graph.nodes[i]->color);
  }

  free_graph(&graph);
  return output;
}