baby_giant_step.cpp

// clang++ -std=c++17 baby_giant_step.cpp -o baby_giant_step -march=native -O3
// ./baby_giant_step

#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <unordered_map>
#include <utility> // std::swap

struct result {
  std::int64_t s;
  std::int64_t t;
  std::int64_t gcd;
};

// Used so that we can use structured bindings nicely
// (i.e. in a ring buffer kind of way).
template <typename T> struct h {
  T &a0;
  T &a1;
  T &a2;

  h(T *arr, std::uint32_t offset)
      : a0{arr[offset % 3]}, a1{arr[(offset + 2) % 3]},
        a2{arr[(offset + 1) % 3]} {}
};

template <typename T> h(T *, std::uint32_t)->h<T>;

result eea(std::int64_t r0, std::int64_t r1) noexcept {
  if (r0 <= r1) {
    std::swap(r0, r1);
  }

  std::int64_t r[3] = {r0, r1, 0};
  std::int64_t s[3] = {1, 0, 0};
  std::int64_t t[3] = {0, 1, 0};

  // Can't declare it in the for loop
  std::uint32_t i = 2;
  for (; r[(i + 2) % 3] != 0; ++i) {
    // This might be oddly named.
    // x0 = x_i, x1 = x_{i-1}, x2 = x_{i-2}
    auto [r0, r1, r2] = h{r, i};
    auto [s0, s1, s2] = h{s, i};
    auto [t0, t1, t2] = h{t, i};

    r0 = r2 % r1;

    const auto q = (r2 - r0) / r1;

    s0 = s2 - q * s1;
    t0 = t2 - q * t1;
  }

  return {s[(i + 1) % 3], t[(i + 1) % 3], r[(i + 1) % 3]};
}

// square multiply
unsigned z_exp(unsigned b, unsigned e, unsigned m) {
  if (e == 0) {
    return 1;
  }

  unsigned acc = b;
  unsigned i = __builtin_clz(e) + 1;
  while (i < 32) {
    acc *= acc;
    acc %= m;
    if (e & (1 << (31 - i++))) {
      acc *= b;
      acc %= m;
    }
  }

  return acc;
}

std::optional<unsigned> baby_giant_step(unsigned a, unsigned b, unsigned p) {
  const auto m = (unsigned)ceil(sqrt(p));
  const auto [_, a_inv, __] = eea(a, p);
  const auto a_inv_m = z_exp(a_inv, m, p);

  // a^{x_b} to x_b
  std::unordered_map<unsigned, unsigned> lookup;
  for (auto i = 0u; i < m; ++i) {
    lookup[z_exp(a, i, p)] = i;
  }

  for (auto i = 0u; i < m; ++i) {
    const auto res = (z_exp(a_inv_m, i, p) * b) % p;
    if (const auto r = lookup.find(res); r != lookup.end()) {
      return {i * m + r->second};
    }
  }

  return {};
}

#include <iostream>

template <typename T>
std::ostream &operator<<(std::ostream &os, const std::optional<T> &x) {
  if (x) {
    return os << *x;
  }
  return os << "(none)";
}

int main(int, char **) {
  std::cout << "a = 17 b = 101: " << baby_giant_step(17, 101, 167)
            << "\na = 17 b = 102: " << baby_giant_step(17, 102, 167)
            << "\na = 18 b = 103: " << baby_giant_step(18, 103, 167)
            << std::endl;

  return 0;
}