defmodule BloomFilter do
  @moduledoc """
  Fixed-size Bloom filter over `:atomics`.
  """

  import Bitwise

  @word_size 64

  defstruct [:ref, :mask, :probes]

  def new(bits, probes) when band(bits, bits - 1) == 0 do
    words = div(bits, @word_size)

    %__MODULE__{
      ref: :atomics.new(words, signed: false),
      mask: bits - 1,
      probes: probes
    }
  end

  def put(%__MODULE__{ref: ref} = filter, term) do
    positions = indexes(filter, term)

    Enum.each(positions, &set_bit(ref, &1))
  end

  def member?(%__MODULE__{ref: ref} = filter, term) do
    positions = indexes(filter, term)

    Enum.all?(positions, &bit_set?(ref, &1))
  end

  defp indexes(%__MODULE__{mask: mask, probes: probes}, term) do
    <<base::64, step::64, _rest::binary>> = :crypto.hash(:sha256, term)

    Enum.map(0..(probes - 1), fn probe ->
      base + probe * step &&& mask
    end)
  end

  defp set_bit(ref, index) do
    word = word_index(index)
    bit = bit_mask(index)

    set_word(ref, word, bit)
  end

  defp set_word(ref, word, bit) do
    current = :atomics.get(ref, word)
    desired = current ||| bit

    if current == desired do
      :ok
    else
      case :atomics.compare_exchange(ref, word, current, desired) do
        :ok -> :ok
        _changed -> set_word(ref, word, bit)
      end
    end
  end

  defp bit_set?(ref, index) do
    word = word_index(index)
    bit = bit_mask(index)
    value = :atomics.get(ref, word)

    (value &&& bit) != 0
  end

  defp word_index(index) do
    div(index, @word_size) + 1
  end

  defp bit_mask(index) do
    1 <<< rem(index, @word_size)
  end
end