libdpf/python/tests/test_pydpf.py

"""Smoke tests for the optional pydpf module."""

import pydpf


def test_point_mass():
    keys = pydpf.make_dpf(7, 11)
    on = pydpf.reconstruct(
        pydpf.eval_point(keys, 0, 7),
        pydpf.eval_point(keys, 1, 7),
    )
    off = pydpf.reconstruct(
        pydpf.eval_point(keys, 0, 8),
        pydpf.eval_point(keys, 1, 8),
    )
    assert on == 11
    assert off == 0


def test_eval_until_prefix_resume():
    keys = pydpf.make_idpf16(0xBEEF)
    ctx0 = pydpf.make_eval_ctx(keys, 0)
    ctx1 = pydpf.make_eval_ctx(keys, 1)
    # High bit of 0xBEEF is 1.
    s0 = pydpf.eval_until(ctx0, 1, [0, 1])
    s1 = pydpf.eval_until(ctx1, 1, [0, 1])
    assert pydpf.ctx_node_count(ctx0) == 2
    assert pydpf.reconstruct(s0[0], s1[0]) == 0
    assert pydpf.reconstruct(s0[1], s1[1]) == 1


def test_it_dpf3_reconstruction():
    keys = pydpf.make_it_dpf3(42, 7)
    mask = (1 << 64) - 1
    on = (
        pydpf.eval_it_dpf3(keys, 0, 42)
        + pydpf.eval_it_dpf3(keys, 1, 42)
        + pydpf.eval_it_dpf3(keys, 2, 42)
    ) & mask
    off = (
        pydpf.eval_it_dpf3(keys, 0, 41)
        + pydpf.eval_it_dpf3(keys, 1, 41)
        + pydpf.eval_it_dpf3(keys, 2, 41)
    ) & mask
    assert on == 7
    assert off == 0


def _open_vec(s0, s1):
    return [pydpf.reconstruct(a, b) for a, b in zip(s0, s1)]


def test_eval_full():
    keys = pydpf.make_dpf(7, 11)
    opened = _open_vec(pydpf.eval_full(keys, 0), pydpf.eval_full(keys, 1))
    assert len(opened) == 256
    assert opened[7] == 11
    assert all(v == 0 for i, v in enumerate(opened) if i != 7)


def test_eval_interval_matches_full_slice():
    keys = pydpf.make_dpf(40, 9)
    full = _open_vec(pydpf.eval_full(keys, 0), pydpf.eval_full(keys, 1))
    lo, hi = 30, 50
    interval = _open_vec(
        pydpf.eval_interval(keys, 0, lo, hi),
        pydpf.eval_interval(keys, 1, lo, hi),
    )
    assert interval == full[lo : hi + 1]


def test_eval_sequence_matches_point():
    keys = pydpf.make_dpf(5, 3)
    points = [1, 5, 5, 200]
    seq = _open_vec(
        pydpf.eval_sequence(keys, 0, points),
        pydpf.eval_sequence(keys, 1, points),
    )
    for p, got in zip(points, seq):
        expect = pydpf.reconstruct(
            pydpf.eval_point(keys, 0, p),
            pydpf.eval_point(keys, 1, p),
        )
        assert got == expect


def test_eval_sequence_recipe():
    keys = pydpf.make_dpf(5, 3)
    points = [1, 5, 200]
    via_seq = _open_vec(
        pydpf.eval_sequence(keys, 0, points),
        pydpf.eval_sequence(keys, 1, points),
    )
    via_recipe = _open_vec(
        pydpf.eval_sequence_recipe(keys, 0, points),
        pydpf.eval_sequence_recipe(keys, 1, points),
    )
    assert via_seq == via_recipe
    try:
        pydpf.eval_sequence(keys, 0, [3, 1])
        assert False, "expected unsorted rejection"
    except ValueError:
        pass


def test_multileaf():
    keys = pydpf.make_dpf_multi(9, 100, 200)
    assert pydpf.reconstruct(
        pydpf.eval_point_leaf(keys, 0, 0, 9),
        pydpf.eval_point_leaf(keys, 1, 0, 9),
    ) == 100
    assert pydpf.reconstruct(
        pydpf.eval_point_leaf(keys, 0, 1, 9),
        pydpf.eval_point_leaf(keys, 1, 1, 9),
    ) == 200
    assert pydpf.reconstruct(
        pydpf.eval_point_leaf(keys, 0, 0, 8),
        pydpf.eval_point_leaf(keys, 1, 0, 8),
    ) == 0
    full0 = _open_vec(
        pydpf.eval_full_leaf(keys, 0, 0),
        pydpf.eval_full_leaf(keys, 1, 0),
    )
    assert full0[9] == 100
    assert full0[0] == 0


def test_wildcard():
    keys = pydpf.make_dpf_wildcard(12)
    try:
        pydpf.eval_point(keys, 0, 12)
        assert False, "expected eval before assign to fail"
    except RuntimeError:
        pass
    pydpf.assign_wildcard(keys, 77)
    assert pydpf.reconstruct(
        pydpf.eval_point(keys, 0, 12),
        pydpf.eval_point(keys, 1, 12),
    ) == 77
    assert pydpf.reconstruct(
        pydpf.eval_point(keys, 0, 13),
        pydpf.eval_point(keys, 1, 13),
    ) == 0
    opened = _open_vec(pydpf.eval_full(keys, 0), pydpf.eval_full(keys, 1))
    assert opened[12] == 77
    assert opened[0] == 0


if __name__ == "__main__":
    test_point_mass()
    test_eval_until_prefix_resume()
    test_it_dpf3_reconstruction()
    test_eval_full()
    test_eval_interval_matches_full_slice()
    test_eval_sequence_matches_point()
    test_eval_sequence_recipe()
    test_multileaf()
    test_wildcard()
    print("ok")