amduat/src/pel_stack/program_dag/program_dag.c

#include "amduat/pel/program_dag.h"

#include "amduat/pel/opreg_kernel.h"
#include "amduat/pel/opreg_kernel_params.h"
#include "amduat/pel/program_dag_desc.h"

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

typedef struct {
  amduat_artifact_t *outputs;
  size_t outputs_len;
} amduat_pel_node_outputs_t;

typedef struct {
  size_t *order;
  const amduat_pel_kernel_op_desc_t **ops;
  amduat_pel_kernel_params_t *params;
} amduat_pel_program_dag_prepared_t;

static void amduat_prepared_reset(amduat_pel_program_dag_prepared_t *prepared) {
  if (prepared == NULL) {
    return;
  }
  prepared->order = NULL;
  prepared->ops = NULL;
  prepared->params = NULL;
}

static void amduat_prepared_free(amduat_pel_program_dag_prepared_t *prepared) {
  if (prepared == NULL) {
    return;
  }
  free(prepared->order);
  free(prepared->ops);
  free(prepared->params);
  amduat_prepared_reset(prepared);
}

static void amduat_set_result(amduat_pel_execution_result_value_t *result,
                              amduat_pel_execution_status_t status,
                              amduat_pel_execution_error_kind_t kind,
                              uint32_t status_code) {
  if (result == NULL) {
    return;
  }
  result->pel1_version = 1;
  result->status = status;
  result->scheme_ref = amduat_pel_program_dag_scheme_ref();
  result->summary.kind = kind;
  result->summary.status_code = status_code;
  result->diagnostics = NULL;
  result->diagnostics_len = 0;
}

static bool amduat_utf8_is_valid(amduat_octets_t value) {
  size_t i = 0;

  while (i < value.len) {
    uint8_t c = value.data[i];
    if (c <= 0x7f) {
      i += 1;
      continue;
    }
    if ((c & 0xe0u) == 0xc0u) {
      if (i + 1 >= value.len) {
        return false;
      }
      if ((value.data[i + 1] & 0xc0u) != 0x80u) {
        return false;
      }
      if (c < 0xc2u) {
        return false;
      }
      i += 2;
      continue;
    }
    if ((c & 0xf0u) == 0xe0u) {
      if (i + 2 >= value.len) {
        return false;
      }
      if ((value.data[i + 1] & 0xc0u) != 0x80u ||
          (value.data[i + 2] & 0xc0u) != 0x80u) {
        return false;
      }
      if (c == 0xe0u && value.data[i + 1] < 0xa0u) {
        return false;
      }
      if (c == 0xedu && value.data[i + 1] >= 0xa0u) {
        return false;
      }
      i += 3;
      continue;
    }
    if ((c & 0xf8u) == 0xf0u) {
      if (i + 3 >= value.len) {
        return false;
      }
      if ((value.data[i + 1] & 0xc0u) != 0x80u ||
          (value.data[i + 2] & 0xc0u) != 0x80u ||
          (value.data[i + 3] & 0xc0u) != 0x80u) {
        return false;
      }
      if (c == 0xf0u && value.data[i + 1] < 0x90u) {
        return false;
      }
      if (c == 0xf4u && value.data[i + 1] >= 0x90u) {
        return false;
      }
      if (c > 0xf4u) {
        return false;
      }
      i += 4;
      continue;
    }
    return false;
  }

  return true;
}

static int amduat_find_node_index(const amduat_pel_program_t *program,
                                  amduat_pel_node_id_t node_id) {
  size_t i;

  for (i = 0; i < program->nodes_len; ++i) {
    if (program->nodes[i].id == node_id) {
      return (int)i;
    }
  }

  return -1;
}

static bool amduat_build_node_order(const amduat_pel_program_t *program,
                                    size_t *out_order) {
  size_t n;
  size_t *deps;
  bool *placed;
  size_t i;

  n = program->nodes_len;
  deps = NULL;
  placed = NULL;

  if (n == 0) {
    return true;
  }

  deps = (size_t *)calloc(n, sizeof(*deps));
  placed = (bool *)calloc(n, sizeof(*placed));
  if (deps == NULL || placed == NULL) {
    free(deps);
    free(placed);
    return false;
  }

  for (i = 0; i < n; ++i) {
    size_t j;
    for (j = i + 1; j < n; ++j) {
      if (program->nodes[i].id == program->nodes[j].id) {
        free(deps);
        free(placed);
        return false;
      }
    }
  }

  for (i = 0; i < n; ++i) {
    size_t j;
    const amduat_pel_node_t *node;

    node = &program->nodes[i];
    for (j = 0; j < node->inputs_len; ++j) {
      const amduat_pel_dag_input_t *input = &node->inputs[j];
      if (input->kind == AMDUAT_PEL_DAG_INPUT_NODE) {
        if (amduat_find_node_index(program, input->value.node.node_id) < 0) {
          free(deps);
          free(placed);
          return false;
        }
        deps[i] += 1;
      } else if (input->kind != AMDUAT_PEL_DAG_INPUT_EXTERNAL) {
        free(deps);
        free(placed);
        return false;
      }
    }
  }

  for (i = 0; i < n; ++i) {
    size_t best = SIZE_MAX;
    size_t j;
    amduat_pel_node_id_t best_id = 0;

    for (j = 0; j < n; ++j) {
      if (placed[j] || deps[j] != 0) {
        continue;
      }
      if (best == SIZE_MAX || program->nodes[j].id < best_id) {
        best = j;
        best_id = program->nodes[j].id;
      }
    }

    if (best == SIZE_MAX) {
      free(deps);
      free(placed);
      return false;
    }

    if (out_order != NULL) {
      out_order[i] = best;
    }
    placed[best] = true;

    for (j = 0; j < n; ++j) {
      size_t k;
      const amduat_pel_node_t *node;

      if (placed[j]) {
        continue;
      }
      node = &program->nodes[j];
      for (k = 0; k < node->inputs_len; ++k) {
        const amduat_pel_dag_input_t *input = &node->inputs[k];
        if (input->kind == AMDUAT_PEL_DAG_INPUT_NODE &&
            input->value.node.node_id == program->nodes[best].id) {
          if (deps[j] == 0) {
            free(deps);
            free(placed);
            return false;
          }
          deps[j] -= 1;
        }
      }
    }
  }

  free(deps);
  free(placed);
  return true;
}

static void amduat_free_node_outputs(amduat_pel_node_outputs_t *node_outputs,
                                     size_t len) {
  size_t i;

  if (node_outputs == NULL) {
    return;
  }

  for (i = 0; i < len; ++i) {
    amduat_pel_node_outputs_t *entry = &node_outputs[i];
    size_t j;
    for (j = 0; j < entry->outputs_len; ++j) {
      free((void *)entry->outputs[j].bytes.data);
      entry->outputs[j].bytes.data = NULL;
      entry->outputs[j].bytes.len = 0;
    }
    free(entry->outputs);
    entry->outputs = NULL;
    entry->outputs_len = 0;
  }
  free(node_outputs);
}

static bool amduat_copy_artifact(amduat_artifact_t *out,
                                 const amduat_artifact_t *src) {
  uint8_t *buffer = NULL;

  if (out == NULL || src == NULL) {
    return false;
  }

  if (src->bytes.len != 0) {
    buffer = (uint8_t *)malloc(src->bytes.len);
    if (buffer == NULL) {
      return false;
    }
    if (src->bytes.data != NULL) {
      memcpy(buffer, src->bytes.data, src->bytes.len);
    }
  }

  out->bytes = amduat_octets(buffer, src->bytes.len);
  out->has_type_tag = src->has_type_tag;
  out->type_tag = src->type_tag;
  return true;
}

static bool amduat_program_prepare(const amduat_pel_program_t *program,
                                   amduat_pel_program_dag_prepared_t *prepared) {
  size_t i;

  if (program == NULL || prepared == NULL) {
    return false;
  }

  amduat_prepared_reset(prepared);

  if (program->nodes_len > 0 && program->nodes == NULL) {
    return false;
  }
  if (program->roots_len > 0 && program->roots == NULL) {
    return false;
  }

  if (program->nodes_len == 0) {
    if (program->roots_len != 0) {
      return false;
    }
    return true;
  }

  prepared->order = (size_t *)malloc(program->nodes_len *
                                     sizeof(*prepared->order));
  prepared->ops = (const amduat_pel_kernel_op_desc_t **)calloc(
      program->nodes_len, sizeof(*prepared->ops));
  prepared->params = (amduat_pel_kernel_params_t *)calloc(
      program->nodes_len, sizeof(*prepared->params));

  if (prepared->order == NULL || prepared->ops == NULL ||
      prepared->params == NULL) {
    amduat_prepared_free(prepared);
    return false;
  }

  for (i = 0; i < program->nodes_len; ++i) {
    const amduat_pel_node_t *node = &program->nodes[i];
    const amduat_pel_kernel_op_desc_t *desc;

    if (node->op.name.len > 0 && node->op.name.data == NULL) {
      amduat_prepared_free(prepared);
      return false;
    }
    if (!amduat_utf8_is_valid(node->op.name)) {
      amduat_prepared_free(prepared);
      return false;
    }
    if (node->inputs_len > 0 && node->inputs == NULL) {
      amduat_prepared_free(prepared);
      return false;
    }
    if (node->params.len > 0 && node->params.data == NULL) {
      amduat_prepared_free(prepared);
      return false;
    }

    desc = amduat_pel_kernel_op_lookup(node->op.name, node->op.version);
    if (desc == NULL) {
      amduat_prepared_free(prepared);
      return false;
    }
    if (node->inputs_len < desc->min_inputs ||
        node->inputs_len > desc->max_inputs) {
      amduat_prepared_free(prepared);
      return false;
    }
    if (!amduat_pel_kernel_params_decode(desc, node->params,
                                         &prepared->params[i])) {
      amduat_prepared_free(prepared);
      return false;
    }

    prepared->ops[i] = desc;
  }

  if (!amduat_build_node_order(program, prepared->order)) {
    amduat_prepared_free(prepared);
    return false;
  }

  for (i = 0; i < program->nodes_len; ++i) {
    const amduat_pel_node_t *node = &program->nodes[i];
    size_t j;
    for (j = 0; j < node->inputs_len; ++j) {
      const amduat_pel_dag_input_t *input = &node->inputs[j];
      if (input->kind == AMDUAT_PEL_DAG_INPUT_NODE) {
        int dep_index = amduat_find_node_index(program,
                                               input->value.node.node_id);
        if (dep_index < 0) {
          amduat_prepared_free(prepared);
          return false;
        }
        if (input->value.node.output_index >=
            prepared->ops[dep_index]->outputs_len) {
          amduat_prepared_free(prepared);
          return false;
        }
      } else if (input->kind != AMDUAT_PEL_DAG_INPUT_EXTERNAL) {
        amduat_prepared_free(prepared);
        return false;
      }
    }
  }

  for (i = 0; i < program->roots_len; ++i) {
    int root_index = amduat_find_node_index(program, program->roots[i].node_id);
    if (root_index < 0) {
      amduat_prepared_free(prepared);
      return false;
    }
    if (program->roots[i].output_index >=
        prepared->ops[root_index]->outputs_len) {
      amduat_prepared_free(prepared);
      return false;
    }
  }

  return true;
}

bool amduat_pel_program_dag_validate(const amduat_pel_program_t *program) {
  amduat_pel_program_dag_prepared_t prepared;
  bool ok;

  amduat_prepared_reset(&prepared);
  ok = amduat_program_prepare(program, &prepared);
  amduat_prepared_free(&prepared);
  return ok;
}

bool amduat_pel_program_dag_exec(
    const amduat_pel_program_t *program,
    const amduat_artifact_t *inputs,
    size_t inputs_len,
    amduat_artifact_t **out_outputs,
    size_t *out_outputs_len,
    amduat_pel_execution_result_value_t *out_result) {
  amduat_pel_program_dag_prepared_t prepared;
  amduat_pel_node_outputs_t *node_outputs;
  amduat_artifact_t *resolved_inputs;
  size_t max_inputs;
  size_t i;

  if (out_outputs == NULL || out_outputs_len == NULL || out_result == NULL) {
    return false;
  }

  *out_outputs = NULL;
  *out_outputs_len = 0;

  amduat_prepared_reset(&prepared);
  if (!amduat_program_prepare(program, &prepared)) {
    amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_INVALID_PROGRAM,
                      AMDUAT_PEL_EXEC_ERROR_PROGRAM, 2);
    return true;
  }

  if (inputs_len > 0 && inputs == NULL) {
    bool needs_inputs = false;
    for (i = 0; i < program->nodes_len; ++i) {
      size_t j;
      const amduat_pel_node_t *node = &program->nodes[i];
      for (j = 0; j < node->inputs_len; ++j) {
        if (node->inputs[j].kind == AMDUAT_PEL_DAG_INPUT_EXTERNAL) {
          needs_inputs = true;
          break;
        }
      }
      if (needs_inputs) {
        break;
      }
    }
    if (needs_inputs) {
      amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_INVALID_INPUTS,
                        AMDUAT_PEL_EXEC_ERROR_INPUTS, 3);
      amduat_prepared_free(&prepared);
      return true;
    }
  }

  if (program->nodes_len == 0) {
    amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_OK,
                      AMDUAT_PEL_EXEC_ERROR_NONE, 0);
    amduat_prepared_free(&prepared);
    return true;
  }

  node_outputs = (amduat_pel_node_outputs_t *)calloc(
      program->nodes_len, sizeof(*node_outputs));
  if (node_outputs == NULL) {
    amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_RUNTIME_FAILED,
                      AMDUAT_PEL_EXEC_ERROR_RUNTIME, 1);
    amduat_prepared_free(&prepared);
    return true;
  }

  max_inputs = 0;
  for (i = 0; i < program->nodes_len; ++i) {
    if (program->nodes[i].inputs_len > max_inputs) {
      max_inputs = program->nodes[i].inputs_len;
    }
  }

  resolved_inputs = NULL;
  if (max_inputs != 0) {
    resolved_inputs = (amduat_artifact_t *)malloc(
        max_inputs * sizeof(*resolved_inputs));
    if (resolved_inputs == NULL) {
      amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_RUNTIME_FAILED,
                        AMDUAT_PEL_EXEC_ERROR_RUNTIME, 1);
      amduat_free_node_outputs(node_outputs, program->nodes_len);
      amduat_prepared_free(&prepared);
      return true;
    }
  }

  for (i = 0; i < program->nodes_len; ++i) {
    size_t node_index = prepared.order[i];
    const amduat_pel_node_t *node = &program->nodes[node_index];
    size_t j;
    uint32_t status_code = 0;

    for (j = 0; j < node->inputs_len; ++j) {
      const amduat_pel_dag_input_t *input = &node->inputs[j];
      if (input->kind == AMDUAT_PEL_DAG_INPUT_EXTERNAL) {
        if (input->value.external.input_index >= inputs_len) {
          amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_INVALID_INPUTS,
                            AMDUAT_PEL_EXEC_ERROR_INPUTS, 3);
          free(resolved_inputs);
          amduat_free_node_outputs(node_outputs, program->nodes_len);
          amduat_prepared_free(&prepared);
          return true;
        }
        resolved_inputs[j] = inputs[input->value.external.input_index];
      } else if (input->kind == AMDUAT_PEL_DAG_INPUT_NODE) {
        int dep_index = amduat_find_node_index(program,
                                               input->value.node.node_id);
        if (dep_index < 0 ||
            input->value.node.output_index >=
                node_outputs[dep_index].outputs_len) {
          amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_INVALID_PROGRAM,
                            AMDUAT_PEL_EXEC_ERROR_PROGRAM, 2);
          free(resolved_inputs);
          amduat_free_node_outputs(node_outputs, program->nodes_len);
          amduat_prepared_free(&prepared);
          return true;
        }
        resolved_inputs[j] =
            node_outputs[dep_index].outputs[input->value.node.output_index];
      } else {
        amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_INVALID_PROGRAM,
                          AMDUAT_PEL_EXEC_ERROR_PROGRAM, 2);
        free(resolved_inputs);
        amduat_free_node_outputs(node_outputs, program->nodes_len);
        amduat_prepared_free(&prepared);
        return true;
      }
    }

    if (!amduat_pel_kernel_op_eval(
            prepared.ops[node_index], resolved_inputs, node->inputs_len,
            &prepared.params[node_index], &node_outputs[node_index].outputs,
            &node_outputs[node_index].outputs_len, &status_code)) {
      if (status_code == 2 || status_code == 3 || status_code == 0) {
        status_code = 1;
      }
      amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_RUNTIME_FAILED,
                        AMDUAT_PEL_EXEC_ERROR_RUNTIME, status_code);
      free(resolved_inputs);
      amduat_free_node_outputs(node_outputs, program->nodes_len);
      amduat_prepared_free(&prepared);
      return true;
    }
  }

  if (program->roots_len != 0) {
    *out_outputs = (amduat_artifact_t *)calloc(
        program->roots_len, sizeof(**out_outputs));
    if (*out_outputs == NULL) {
      amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_RUNTIME_FAILED,
                        AMDUAT_PEL_EXEC_ERROR_RUNTIME, 1);
      free(resolved_inputs);
      amduat_free_node_outputs(node_outputs, program->nodes_len);
      amduat_prepared_free(&prepared);
      return true;
    }
  }

  for (i = 0; i < program->roots_len; ++i) {
    int root_index = amduat_find_node_index(program, program->roots[i].node_id);
    if (root_index < 0 ||
        program->roots[i].output_index >=
            node_outputs[root_index].outputs_len) {
      amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_INVALID_PROGRAM,
                        AMDUAT_PEL_EXEC_ERROR_PROGRAM, 2);
      amduat_pel_program_dag_free_outputs(*out_outputs,
                                          program->roots_len);
      *out_outputs = NULL;
      *out_outputs_len = 0;
      free(resolved_inputs);
      amduat_free_node_outputs(node_outputs, program->nodes_len);
      amduat_prepared_free(&prepared);
      return true;
    }
    if (!amduat_copy_artifact(
            &(*out_outputs)[i],
            &node_outputs[root_index]
                 .outputs[program->roots[i].output_index])) {
      amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_RUNTIME_FAILED,
                        AMDUAT_PEL_EXEC_ERROR_RUNTIME, 1);
      amduat_pel_program_dag_free_outputs(*out_outputs,
                                          program->roots_len);
      *out_outputs = NULL;
      *out_outputs_len = 0;
      free(resolved_inputs);
      amduat_free_node_outputs(node_outputs, program->nodes_len);
      amduat_prepared_free(&prepared);
      return true;
    }
  }

  *out_outputs_len = program->roots_len;
  amduat_set_result(out_result, AMDUAT_PEL_EXEC_STATUS_OK,
                    AMDUAT_PEL_EXEC_ERROR_NONE, 0);

  free(resolved_inputs);
  amduat_free_node_outputs(node_outputs, program->nodes_len);
  amduat_prepared_free(&prepared);
  return true;
}

void amduat_pel_program_dag_free_outputs(amduat_artifact_t *outputs,
                                         size_t outputs_len) {
  size_t i;

  if (outputs == NULL) {
    return;
  }

  for (i = 0; i < outputs_len; ++i) {
    free((void *)outputs[i].bytes.data);
    outputs[i].bytes.data = NULL;
    outputs[i].bytes.len = 0;
  }
  free(outputs);
}