Modify discrete log to support sync and async (#405)

* Modify discrete log to support sync and async

* Address PR feedback

* update version

Author: AddressXception

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "electionguard"
-version = "1.2.2"
+version = "1.2.3"
 description = "ElectionGuard: Support for e2e verified elections."
 license = "MIT"
 authors = ["Microsoft <electionguard@microsoft.com>"]

src/electionguard/decrypt_with_shares.py

@@ -6,7 +6,7 @@
     DecryptionShare,
     get_shares_for_selection,
 )
-from .dlog import discrete_log
+from .discrete_log import DiscreteLog
 from .group import ElementModP, ElementModQ, mult_p, div_p
 from .tally import (
     CiphertextTally,
@@ -165,7 +165,7 @@ def decrypt_selection_with_decryption_shares(
 
     # Calculate 𝑀=𝐵⁄(∏𝑀𝑖) mod 𝑝.
     decrypted_value = div_p(selection.ciphertext.data, all_shares_product_M)
-    d_log = discrete_log(decrypted_value)
+    d_log = DiscreteLog().discrete_log(decrypted_value)
     return PlaintextTallySelection(
         selection.object_id,
         d_log,

src/electionguard/discrete_log.py

@@ -0,0 +1,101 @@
+# pylint: disable=global-statement
+# support for computing discrete logs, with a cache so they're never recomputed
+
+import asyncio
+from typing import Dict, Tuple
+
+from electionguard.singleton import Singleton
+
+from .group import G, ElementModP, ONE_MOD_P, mult_p, int_to_p_unchecked
+
+DLOG_CACHE = Dict[ElementModP, int]
+DLOG_MAX = 100_000_000
+"""The max number to calculate.  This value is used to stop a race condition."""
+
+
+def discrete_log(element: ElementModP, cache: DLOG_CACHE) -> Tuple[int, DLOG_CACHE]:
+    """
+    Computes the discrete log (base g, mod p) of the given element,
+    with internal caching of results. Should run efficiently when called
+    multiple times when the exponent is at most in the single-digit millions.
+    Performance will degrade if it's much larger.
+
+    For the best possible performance,
+    pre-compute the discrete log of a number you expect to have the biggest
+    exponent you'll ever see. After that, the cache will be fully loaded,
+    and every call will be nothing more than a dictionary lookup.
+    """
+
+    if element in cache:
+        return (cache[element], cache)
+
+    cache = compute_discrete_log_cache(element, cache)
+    return (cache[element], cache)
+
+
+async def discrete_log_async(
+    element: ElementModP,
+    cache: DLOG_CACHE,
+    mutex: asyncio.Lock = asyncio.Lock(),
+) -> Tuple[int, DLOG_CACHE]:
+    """
+    Computes the discrete log (base g, mod p) of the given element,
+    with internal caching of results. Should run efficiently when called
+    multiple times when the exponent is at most in the single-digit millions.
+    Performance will degrade if it's much larger.
+
+    Note: *this function is thread-safe*. For the best possible performance,
+    pre-compute the discrete log of a number you expect to have the biggest
+    exponent you'll ever see. After that, the cache will be fully loaded,
+    and every call will be nothing more than a dictionary lookup.
+    """
+    if element in cache:
+        return (cache[element], cache)
+
+    async with mutex:
+        if element in cache:
+            return (cache[element], cache)
+
+        cache = compute_discrete_log_cache(element, cache)
+        return (cache[element], cache)
+
+
+def compute_discrete_log_cache(
+    element: ElementModP, cache: DLOG_CACHE
+) -> Dict[ElementModP, int]:
+    """
+    Compute a discrete log cache up to the specified element.
+    """
+    if not cache:
+        cache = {ONE_MOD_P: 0}
+    max_element = list(cache)[-1]
+    exponent = cache[max_element]
+
+    g = int_to_p_unchecked(G)
+    while element != max_element:
+        exponent = exponent + 1
+        if exponent > DLOG_MAX:
+            raise ValueError("size is larger than max.")
+        max_element = mult_p(g, max_element)
+        cache[max_element] = exponent
+        print(f"max: {max_element}, exp: {exponent}")
+    return cache
+
+
+class DiscreteLog(Singleton):
+    """
+    A class instance of the discrete log that includes a cache.
+    """
+
+    _cache: DLOG_CACHE = {ONE_MOD_P: 0}
+    _mutex = asyncio.Lock()
+
+    def discrete_log(self, element: ElementModP) -> int:
+        (result, cache) = discrete_log(element, self._cache)
+        self._cache = cache
+        return result
+
+    async def discrete_log_async(self, element: ElementModP) -> int:
+        (result, cache) = await discrete_log_async(element, self._cache, self._mutex)
+        self._cache = cache
+        return result

src/electionguard/dlog.py

@@ -1,6 +1,6 @@
 from typing import Iterable, NamedTuple, Optional
 
-from .dlog import discrete_log
+from .discrete_log import DiscreteLog
 from .group import (
     ElementModQ,
     ElementModP,
@@ -48,7 +48,7 @@ def decrypt_known_product(self, product: ElementModP) -> int:
         :param product: The known product (blinding factor).
         :return: An exponentially encoded plaintext message.
         """
-        return discrete_log(mult_p(self.data, mult_inv_p(product)))
+        return DiscreteLog().discrete_log(mult_p(self.data, mult_inv_p(product)))
 
     def decrypt(self, secret_key: ElementModQ) -> int:
         """

src/electionguard/elgamal.py

@@ -0,0 +1,101 @@
+import unittest
+
+from hypothesis import given
+from hypothesis.strategies import integers
+
+from electionguard.discrete_log import (
+    discrete_log,
+    discrete_log_async,
+    DiscreteLog,
+)
+from electionguard.group import (
+    ElementModP,
+    ONE_MOD_P,
+    mult_p,
+    G,
+    g_pow_p,
+    int_to_q,
+    int_to_p_unchecked,
+    int_to_q_unchecked,
+    P,
+)
+from electionguard.utils import get_optional
+
+
+def _discrete_log_uncached(e: ElementModP) -> int:
+    """
+    A simpler implementation of discrete_log, only meant for comparison testing of the caching version.
+    """
+    count = 0
+    g_inv = int_to_p_unchecked(pow(G, -1, P))
+    while e != ONE_MOD_P:
+        e = mult_p(e, g_inv)
+        count = count + 1
+
+    return count
+
+
+class TestDiscreteLogFunctions(unittest.TestCase):
+    """Discrete log tests"""
+
+    @given(integers(0, 100))
+    def test_uncached(self, exp: int):
+        plaintext = get_optional(int_to_q(exp))
+        exp_plaintext = g_pow_p(plaintext)
+        plaintext_again = _discrete_log_uncached(exp_plaintext)
+
+        self.assertEqual(exp, plaintext_again)
+
+    @given(integers(0, 1000))
+    def test_cached(self, exp: int):
+        cache = {ONE_MOD_P: 0}
+        plaintext = get_optional(int_to_q(exp))
+        exp_plaintext = g_pow_p(plaintext)
+        (plaintext_again, returned_cache) = discrete_log(exp_plaintext, cache)
+
+        self.assertEqual(exp, plaintext_again)
+        self.assertEqual(len(cache), len(returned_cache))
+
+    def test_cached_one(self):
+        cache = {ONE_MOD_P: 0}
+        plaintext = int_to_q_unchecked(1)
+        ciphertext = g_pow_p(plaintext)
+        (plaintext_again, returned_cache) = discrete_log(ciphertext, cache)
+
+        self.assertEqual(1, plaintext_again)
+        self.assertEqual(len(cache), len(returned_cache))
+
+    async def test_cached_one_async(self):
+        cache = {ONE_MOD_P: 0}
+        plaintext = int_to_q_unchecked(1)
+        ciphertext = g_pow_p(plaintext)
+        (plaintext_again, returned_cache) = await discrete_log_async(ciphertext, cache)
+
+        self.assertEqual(1, plaintext_again)
+        self.assertEqual(len(cache), len(returned_cache))
+
+
+class TestDiscreteLogClass(unittest.TestCase):
+    """Discrete log tests"""
+
+    @given(integers(0, 1000))
+    def test_cached(self, exp: int):
+        plaintext = get_optional(int_to_q(exp))
+        exp_plaintext = g_pow_p(plaintext)
+        plaintext_again = DiscreteLog().discrete_log(exp_plaintext)
+
+        self.assertEqual(exp, plaintext_again)
+
+    def test_cached_one(self):
+        plaintext = int_to_q_unchecked(1)
+        ciphertext = g_pow_p(plaintext)
+        plaintext_again = DiscreteLog().discrete_log(ciphertext)
+
+        self.assertEqual(1, plaintext_again)
+
+    async def test_cached_one_async(self):
+        plaintext = int_to_q_unchecked(1)
+        ciphertext = g_pow_p(plaintext)
+        plaintext_again = await DiscreteLog().discrete_log_async(ciphertext)
+
+        self.assertEqual(1, plaintext_again)