proto_impl.c

#include "proto.h"
#include <assert.h>
#include <stdlib.h>

#include "vendor/libtommath/tommath.h"

_Static_assert(MP_DIGIT_BIT == 60, "");
_Static_assert(sizeof(mp_digit) == 8, "");

_Static_assert(offsetof(struct proto_ty, digits) == 24, "");

static bool can_alloc(proto_context ctx, size_t bytes)
{
  (void)bytes;
  uint64_t s = proto_context_fuel_value(ctx);
  if (s == 0)
    {
      // fprintf(stderr, "Out of memory tokens!\n");
      return false;
    }

  proto_context_fuel_decrement(ctx);
  proto_context_count_increment(ctx);
  return true;
}

// Cast carefully as it's the same underlying memory
static mp_int proto_to_mp_int(proto x)
{
  mp_int res;
  res.used = proto_used(x);
  res.alloc = proto_alloced(x);
  res.sign = proto_zpos(x) ? MP_ZPOS : MP_NEG;
  res.dp = &x->digits[0];
  return res;
}

static proto mp_int_to_proto(mp_int *x)
{
  // The mp_int instance may have updated used/zpos fields which
  // aren't yet reflected in the underlying proto instance
  proto res = (proto)((char *)x->dp - offsetof(struct proto_ty, digits));
  res->used = x->used;
  res->zpos = (x->sign == MP_ZPOS);
  return res;
}

proto proto_create_default_size(proto_context ctx)
{
  // docs say the alloc size should never be zero
  // mp_init passes MP_PREC which is probably 8
  // current tests pass with zero though, lets try it
  if (proto_context_fuel_value(ctx) != UINT64_MAX)
    return proto_create(ctx, 0);
  else
    return proto_create(ctx, 1); // todo should probably be the same as mp_init
}


proto proto_create(proto_context ctx, size_t digits)
{  
  size_t bytes = sizeof(struct proto_ty) + digits * sizeof(uint64_t);
  proto state = can_alloc(ctx, bytes) ? malloc(bytes) : 0;
  if (state)
    {
      state->used = 0;
      state->alloc = digits;
      state->zpos = true;
      for (size_t i = 0; i < digits; i++)
        {
          state->digits[i] = 0;
        }
    }
  return state;
}

proto proto_create_invalid(void) { return NULL; }

bool proto_valid(proto_context ctx, proto x) { return x != NULL; }

void proto_destroy(proto_context ctx, proto x) { free(x); }

proto proto_copy(proto_context ctx, proto x)
{
  size_t digits = proto_alloced(x);
  proto res = proto_create(ctx, digits);
  if (res)
    {
      res->used = x->used;
      res->alloc = x->alloc;
      res->zpos = x->zpos;
      for (size_t i = 0; i < digits; i++)
        {
          res->digits[i] = x->digits[i];
        }
    }
  return res;
}

void proto_dump(proto_context ctx, proto x)
{
  if (!proto_valid(ctx, x))
    {
      printf("(proto) <invalid>\n");
      return;
    }
  if (proto_is_sentinel(x))
    {
      printf("(proto) <sentinel>\n");
      return;
    }
  size_t alloc = x->alloc;
  printf(
      "(proto)\n"
      "{\n"
      "  .used = %zu,\n"
      "  .alloc = %zu,\n"
      "  .zpos = %u,\n"
      "  .digits[%lu] =\n  {\n",
      x->used, alloc, x->zpos, alloc);

  for (size_t i = 0; i < alloc; i++)
    {
      printf("    %lu,\n", x->digits[i]);
    }

  printf(
      ""
      "  },\n"
      "}\n");
}

proto proto_sentinel(void)
{
  static struct proto_ty p = {
      .used = ~0,
      .alloc = ~0,
      .zpos = false,
  };
  return &p;
}

bool proto_is_sentinel(proto x) { return x == proto_sentinel(); }

size_t proto_used(proto x) { return x->used; }
size_t proto_alloced(proto x) { return x->alloc; }
bool proto_zpos(proto x) { return x->zpos; }

bool proto_is_zero(proto_context ctx, proto x)
{
  uint64_t N = proto_used(x);
  for (uint64_t i = 0; i < N; i++)
    {
      if (x->digits[i] != 0)
        {
          return false;
        }
    }
  return true;
}

bool proto_resize(proto_context ctx, proto *x, size_t digits)
{
  size_t alloced = proto_alloced(*x);
  size_t bytes = sizeof(struct proto_ty) + digits * sizeof(uint64_t);
  proto p = *x;
  proto r = can_alloc(ctx, bytes) ? realloc(p, bytes) : 0;
  if (r)
    {
      r->alloc = digits;
      for (size_t i = alloced; i < digits; i++)
        {
          r->digits[i] = 0;
        }
      *x = r;
      return true;
    }
  else
    {
      return false;
    }
}

bool proto_equal(proto_context ctx, proto x, proto y)
{
  uint64_t N = proto_used(x);
  if (N != proto_used(y))
    {
      return false;
    }

  if (x->zpos != y->zpos)
    {
      return false;
    }

  for (uint64_t i = 0; i < N; i++)
    {
      if (x->digits[i] != y->digits[i])
        {
          return false;
        }
    }

  return true;
}

// Implement the rest in terms of libtommath and careful casting

// Pass a context object around without modifying the libtommath api
static proto_context *global(void);
static void global_set(proto_context *ctx);
static void global_clear(void);

proto proto_from_u32(proto_context ctx, uint32_t val)
{
  proto p = proto_create_default_size(ctx);
  if (!proto_valid(ctx, p))
    {
      return p;
    }

  mp_int tmp = proto_to_mp_int(p);

  global_set(&ctx);
  mp_set_u32(&tmp, val);
  global_clear();

  return mp_int_to_proto(&tmp);
}

uint32_t proto_to_u32(proto_context ctx, proto x)
{
  assert(proto_valid(ctx, x));

  mp_int tmp = proto_to_mp_int(x);
  global_set(&ctx);
  uint32_t res = mp_get_u32(&tmp);
  global_clear();
  return res;
}

static proto proto_unary(proto_context ctx, proto x,
                         mp_err (*func)(const mp_int *, mp_int *))
{
  if (!proto_valid(ctx, x))
    {
      return proto_create_invalid();
    }

  proto y = proto_create_default_size(ctx);
  if (!proto_valid(ctx, y))
    {
      return proto_create_invalid();
    }

  mp_int mx = proto_to_mp_int(x);
  mp_int my = proto_to_mp_int(y);

  global_set(&ctx);
  mp_err res = func(&mx, &my);
  global_clear();

  if (res == MP_OKAY)
    {
      return mp_int_to_proto(&my);
    }
  else
    {
      proto_destroy(ctx, y);
      return proto_create_invalid();
    }
}

static proto proto_binary(proto_context ctx, proto x, proto y,
                          mp_err (*func)(const mp_int *, const mp_int *,
                                         mp_int *))
{
  if (!proto_valid(ctx, x) || !proto_valid(ctx, y))
    {
      return proto_create_invalid();
    }

  proto z = proto_create_default_size(ctx);
  if (!proto_valid(ctx, z))
    {
      return proto_create_invalid();
    }

  mp_int mx = proto_to_mp_int(x);
  mp_int my = proto_to_mp_int(y);
  mp_int mz = proto_to_mp_int(z);

  global_set(&ctx);
  mp_err res = func(&mx, &my, &mz);
  global_clear();

  if (res == MP_OKAY)
    {
      return mp_int_to_proto(&mz);
    }
  else
    {
      proto_destroy(ctx, z);
      return proto_create_invalid();
    }
}

static proto proto_mutating_unary(proto_context ctx, proto x,
                                  mp_err (*func)(mp_int *))
{
  if (!proto_valid(ctx, x))
    {
      return proto_create_invalid();
    }

  proto res = proto_copy(ctx, x);
  if (!proto_valid(ctx, res))
    {
      return proto_create_invalid();
    }

  mp_int mr = proto_to_mp_int(res);
  global_set(&ctx);
  mp_err err = func(&mr);
  global_clear();
  if (err == MP_OKAY)
    {
      return mp_int_to_proto(&mr);
    }
  else
    {
      proto_destroy(ctx, res);
      return proto_create_invalid();
    }
}

static proto proto_op_u32(proto_context ctx, proto x, uint32_t y,
                          mp_err (*func)(const mp_int *, mp_digit, mp_int *))
{
  // This is on thin ice if the underlying size is int32_t on portability grounds
  _Static_assert(MP_DIGIT_MAX >= UINT32_MAX, "");

  mp_int mx = proto_to_mp_int(x);

  proto z = proto_create_default_size(ctx);
  if (!proto_valid(ctx, z))
    {
      return proto_create_invalid();
    }
  mp_int mz = proto_to_mp_int(z);

  global_set(&ctx);
  mp_err err = func(&mx, y, &mz);
  global_clear();

  if (err == MP_OKAY)
    {
      return mp_int_to_proto(&mz);
    }
  else
    {
      proto_destroy(ctx, z);
      return proto_create_invalid();
    }
}

proto proto_abs(proto_context ctx, proto x)
{
  return proto_unary(ctx, x, mp_abs);
}
proto proto_neg(proto_context ctx, proto x)
{
  return proto_unary(ctx, x, mp_neg);
}

proto proto_incr(proto_context ctx, proto x)
{
  return proto_mutating_unary(ctx, x, mp_incr);
}

proto proto_decr(proto_context ctx, proto x)
{
  return proto_mutating_unary(ctx, x, mp_decr);
}

proto proto_add(proto_context ctx, proto x, proto y)
{
  return proto_binary(ctx, x, y, mp_add);
}
proto proto_sub(proto_context ctx, proto x, proto y)
{
  return proto_binary(ctx, x, y, mp_sub);
}
proto proto_mul(proto_context ctx, proto x, proto y)
{
  return proto_binary(ctx, x, y, mp_mul);
}

static mp_err mp_div_quotient(const mp_int *a, const mp_int *b, mp_int *c)
{
  mp_int tmp;
  mp_err err;
  if ((err = mp_init(&tmp)) != MP_OKAY) return err;
  err = mp_div(a, b, c, &tmp);
  mp_clear(&tmp);
  return err;
}

proto proto_div(proto_context ctx, proto x, proto y)
{
  return proto_binary(ctx, x, y, mp_div_quotient);
}

static mp_err mp_div_remainder(const mp_int *a, const mp_int *b, mp_int *c)
{
  mp_int tmp;
  mp_err err;
  if ((err = mp_init(&tmp)) != MP_OKAY) return err;
  err = mp_div(a, b, &tmp, c);
  mp_clear(&tmp);
  return err;
}

proto proto_rem(proto_context ctx, proto x, proto y)
{
  return proto_binary(ctx, x, y, mp_div_remainder);
}

proto proto_not(proto_context ctx, proto x)
{
  return proto_unary(ctx, x, mp_complement);
}
proto proto_or(proto_context ctx, proto x, proto y)
{
  return proto_binary(ctx, x, y, mp_or);
}
proto proto_and(proto_context ctx, proto x, proto y)
{
  return proto_binary(ctx, x, y, mp_and);
}
proto proto_xor(proto_context ctx, proto x, proto y)
{
  return proto_binary(ctx, x, y, mp_xor);
}

// arithmetic shift is called signed_rsh

bool proto_within_int32(proto_context ctx, proto x)
{
  
  if (proto_zpos(x))
    {
      return proto_cmp_enum_u32(ctx, x, INT32_MAX) == proto_cmp_res_lt;
    }
  else
    {
      if (proto_used(x) > 1)
        {
          return false;
        }
      else
        {
          return x->digits[0] < INT32_MAX; // might be off by one here
        }
    }
}

proto proto_ash(proto_context ctx, proto x, proto y)
{
  if (!proto_within_int32(ctx, y))
    {
      return proto_create_invalid();
    }

  // libtommath doesn't ignore negative shifts here, but
  // also doesn't do any obvious handling for them. Just
  // passes into divide (as large unsigned values)
  // but as lsh and rsh are currently returning invalid, do so here too
  if (!proto_zpos(y))
    {
      return proto_create_invalid();
    }

  uint32_t by = proto_to_u32(ctx, y);
  if (by > INT32_MAX) {
    // internal error in within_int32 test
    __builtin_trap();
  }

  proto z = proto_create_default_size(ctx);
  if (!proto_valid(ctx, z))
    {
      return proto_create_invalid();
    }

  mp_int mx = proto_to_mp_int(x);
  mp_int mz = proto_to_mp_int(z);
  
  global_set(&ctx);
  mp_err err = mp_signed_rsh(&mx, (int)by, &mz);
  global_clear();
  if (err == MP_OKAY)
    {
      return mp_int_to_proto(&mz);
    }
  else
    {
      proto_destroy(ctx, z);
      return proto_create_invalid();
    }
}

proto proto_rsh(proto_context ctx, proto x, proto y)
{
  if (!proto_within_int32(ctx, y))
    {
      return proto_create_invalid();
    }

  // left shift by negative amount means what?
  // libtommath thinks it's a no-op, C thinks it's UB
  // reject for now
  if (!proto_zpos(y))
    {
      return proto_create_invalid();
    }

  uint32_t by = proto_to_u32(ctx, y);
  if (by > INT32_MAX) {
    // internal error in within_int32 test
    __builtin_trap();
  }

  proto res = proto_copy(ctx, x);
  if (!proto_valid(ctx, res))
    {
      return proto_create_invalid();
    }

  mp_int mr = proto_to_mp_int(res);
  global_set(&ctx);
  // right shift doesn't fail
  mp_rshd(&mr, (int)by);
  global_clear();

  return mp_int_to_proto(&mr);
}


proto proto_lsh(proto_context ctx, proto x, proto y)
{
  if (!proto_within_int32(ctx, y))
    {
      return proto_create_invalid();
    }

  // left shift by negative amount means what?
  // libtommath thinks it's a no-op, C thinks it's UB
  // reject for now
  if (!proto_zpos(y))
    {
      return proto_create_invalid();
    }

  uint32_t by = proto_to_u32(ctx, y);
  if (by > INT32_MAX) {
    // internal error in within_int32 test
    __builtin_trap();
  }

  proto res = proto_copy(ctx, x);
  if (!proto_valid(ctx, res))
    {
      return proto_create_invalid();
    }

  
  mp_int mr = proto_to_mp_int(res);
  global_set(&ctx);
  mp_err err = mp_lshd(&mr, (int)by);
  global_clear();
  if (err == MP_OKAY)
    {
      return mp_int_to_proto(&mr);
    }
  else
    {
      proto_destroy(ctx, res);
      return proto_create_invalid();
    }  
}


static proto_cmp_res mp_ord_to_proto_cmp(mp_ord x)
{
  switch (x)
    {
    case MP_LT: return proto_cmp_res_lt;
    case MP_EQ: return proto_cmp_res_eq;
    case MP_GT: return proto_cmp_res_gt;
    default:
      return (proto_cmp_res)x;
    }                      

}
proto_cmp_res proto_cmp_enum(proto_context ctx, proto x, proto y)
{
  mp_int mx = proto_to_mp_int(x);
  mp_int my = proto_to_mp_int(y);
  global_set(&ctx);
  mp_ord res = mp_cmp(&mx, &my);
  global_clear();
  return mp_ord_to_proto_cmp(res);
}

proto_cmp_res proto_cmp_enum_u32(proto_context ctx, proto x, uint32_t y)
{
  _Static_assert(MP_DIGIT_MAX >= UINT32_MAX, ""); // dubious
  // if (y < MP_DIGIT_MAX)
    {
      mp_int mx = proto_to_mp_int(x);
      global_set(&ctx);
      mp_ord res = mp_cmp_d(&mx, y);
      global_clear();
      return mp_ord_to_proto_cmp(res);
    }

    __builtin_trap();
}

proto proto_add_u32(proto_context ctx, proto x, uint32_t y)
{
  return proto_op_u32(ctx, x, y, mp_add_d);
}

proto proto_mul_u32(proto_context ctx, proto x, uint32_t y)
{
  return proto_op_u32(ctx, x, y, mp_mul_d);
}

// Care needed with error handling here. Assumes that the functions called
// have no failure modes when passed the same value for both arguments

static proto proto_op_unary_move(proto_context ctx, proto x,
                                 mp_err (*func)(const mp_int *, mp_int *))
{
  if (!proto_valid(ctx, x))
    {
      return proto_create_invalid();
    }

  mp_int mx = proto_to_mp_int(x);

  global_set(&ctx);
  mp_err res = func(&mx, &mx);
  (void)res;  // these do not alloc when passed same value twice
  global_clear();

  return mp_int_to_proto(&mx);
}

static proto proto_op_u32_move(proto_context ctx, proto x, uint32_t y,
                               mp_err (*func)(const mp_int *, mp_digit,
                                              mp_int *))
{
  _Static_assert(MP_DIGIT_MAX >= UINT32_MAX, "");

  mp_int mx = proto_to_mp_int(x);

  global_set(&ctx);
  mp_err err = func(&mx, y, &mx);
  global_clear();

  // These may need to allocate
  // Could lift the c->alloc < (a->used + 1) check
  // They work by making sure there's a word available then clamp at the end
  if (err == MP_OKAY)
    {
      return mp_int_to_proto(&mx);
    }
  else
    {
      proto_destroy(ctx, x);
      return proto_create_invalid();
    }
}

proto proto_abs_move(proto_context ctx, proto x)
{
  return proto_op_unary_move(ctx, x, mp_abs);
}

proto proto_neg_move(proto_context ctx, proto x)
{
  return proto_op_unary_move(ctx, x, mp_neg);
}

proto proto_add_u32_move(proto_context ctx, proto x, uint32_t y)
{
  return proto_op_u32_move(ctx, x, y, mp_add_d);
}

proto proto_mul_u32_move(proto_context ctx, proto x, uint32_t y)
{
  return proto_op_u32_move(ctx, x, y, mp_mul_d);
}

uint64_t proto_digit_max(proto_context ctx)
{
  (void)ctx;
  // hardcoding these for the moment
  _Static_assert(MP_DIGIT_BIT  ==  60, "");
  _Static_assert(MP_DIGIT_MAX == 1152921504606846976-1, ""); // 2**60 -1

  // 18 base 10 digitis is the limit
  // (2**60-1) > 10**18 // true
  // (2**60-1) > 10**19 // false

  
  return MP_DIGIT_MAX;
}

uint32_t proto_base_ten_per_digit(proto_context ctx)
{
  (void)ctx;
  _Static_assert(MP_DIGIT_BIT  ==  60, "");
  return 18;
}

// Define the libtommath memory interface non-invasively
// mp_init passes MP_PREC as the size field, i.e. mp_init
// can fail.
mp_err mp_init_size(mp_int *a, int size)
{
  if (size < 0)
    {
      return MP_VAL;
    }

  enum
  {
    sizeof_bits = ((size_t)CHAR_BIT * sizeof(long long)),
    min_prec = ((((int)sizeof_bits + MP_DIGIT_BIT) - 1) / MP_DIGIT_BIT)
  };

  size = (min_prec > size) ? min_prec : size;

  proto_context *ctx = global();
  proto p = proto_create(*ctx, size);
  if (proto_valid(*ctx, p))
    {
      *a = proto_to_mp_int(p);
      return MP_OKAY;
    }
  else
    {
      return MP_MEM;
    }
}

void mp_clear(mp_int *a)
{
  if (a->dp != NULL)
    {
      proto p = mp_int_to_proto(a);
      proto_context *ctx = global();
      proto_destroy(*ctx, p);
      a->dp = NULL;
      a->alloc = 0;
      a->used = 0;
      a->sign = MP_ZPOS;
    }
}


// This is called in a context where 'size' is known to be the designed value
// e.g.
/*
int used = MP_MAX(a->used, b->used) + 1;
if (c->alloc < used) {
  if ((err = mp_grow(c, used)) != MP_OKAY) {
    return err;
  }
}
*/
// This is difficult to catch from a fallable malloc. Used is the number of limbs needed
// in the output, if the value passed in already has some memory allocated, this may never
// trigger
// I think there's an invariant in libtommath that alloc is never zero, as well as tending
// to allocate at least a few words by default


mp_err mp_grow(mp_int *a, int size)
{
  if (size < 0) return MP_VAL;

  proto p = mp_int_to_proto(a);
  proto_context *ctx = global();
  if (proto_resize(*ctx, &p, size))
    {
      *a = proto_to_mp_int(p);
      return MP_OKAY;
    }
  else
    {
      return MP_MEM;
    }
}

#if 0
// None of the mp functions call shrink
mp_err mp_shrink(mp_int *a)
{
  return MP_MEM;
}
#endif

static _Thread_local proto_context *global_mp_ctx = 0;

static proto_context *global(void)
{
  assert(global_mp_ctx != 0);
  return global_mp_ctx;
}
static void global_set(proto_context *ctx)
{
  assert(ctx != 0);
  assert(global_mp_ctx == 0);
  global_mp_ctx = ctx;
}
static void global_clear(void)
{
  assert(global_mp_ctx != 0);
  global_mp_ctx = 0;
}