Add base32_prefix_bits function to get prefix of DAG hash as an int.

2 files changed, 61 insertions, 4 deletions

diff --git a/lib/spack/spack/spec.py b/lib/spack/spack/spec.py
index ba9cea876d..94ac8788ba 100644
--- a/lib/spack/spack/spec.py
+++ b/lib/spack/spack/spec.py
@@ -963,13 +963,10 @@ class Spec(object):
         return Prefix(spack.install_layout.path_for_spec(self))
 
     def dag_hash(self, length=None):
-        """
-        Return a hash of the entire spec DAG, including connectivity.
-        """
+        """Return a hash of the entire spec DAG, including connectivity."""
         if self._hash:
             return self._hash[:length]
         else:
-            # XXX(deptype): ignore 'build' dependencies here
             yaml_text = syaml.dump(
                 self.to_node_dict(), default_flow_style=True, width=sys.maxint)
             sha = hashlib.sha1(yaml_text)
@@ -978,6 +975,10 @@ class Spec(object):
                 self._hash = b32_hash
             return b32_hash
 
+    def dag_hash_bit_prefix(self, bits):
+        """Get the first <bits> bits of the DAG hash as an integer type."""
+        return base32_prefix_bits(self.dag_hash(), bits)
+
     def to_node_dict(self):
         d = syaml_dict()
 
@@ -999,6 +1000,8 @@ class Spec(object):
         if self.architecture:
             d['arch'] = self.architecture.to_dict()
 
+        # TODO: restore build dependencies here once we have less picky
+        # TODO: concretization.
         deps = self.dependencies_dict(deptype=('link', 'run'))
         if deps:
             d['dependencies'] = syaml_dict([
@@ -2723,6 +2726,26 @@ def parse_anonymous_spec(spec_like, pkg_name):
     return anon_spec
 
 
+def base32_prefix_bits(hash_string, bits):
+    """Return the first <bits> bits of a base32 string as an integer."""
+    if bits > len(hash_string) * 5:
+        raise ValueError("Too many bits! Requested %d bit prefix of '%s'."
+                         % (bits, hash_string))
+
+    hash_bytes = base64.b32decode(hash_string, casefold=True)
+
+    result = 0
+    n = 0
+    for i, b in enumerate(hash_bytes):
+        n += 8
+        result = (result << 8) | ord(b)
+        if n >= bits:
+            break
+
+    result >>= (n - bits)
+    return result
+
+
 class SpecError(spack.error.SpackError):
     """Superclass for all errors that occur while constructing specs."""
 
diff --git a/lib/spack/spack/test/spec_dag.py b/lib/spack/spack/test/spec_dag.py
index 40cdb02966..0bc63bcf0f 100644
--- a/lib/spack/spack/test/spec_dag.py
+++ b/lib/spack/spack/test/spec_dag.py
@@ -523,3 +523,37 @@ class SpecDagTest(MockPackagesTest):
             level = descend_and_check(dag.to_node_dict())
             # level just makes sure we are doing something here
             self.assertTrue(level >= 5)
+
+    def test_hash_bits(self):
+        """Ensure getting first n bits of a base32-encoded DAG hash works."""
+
+        # RFC 4648 base32 decode table
+        b32 = dict((j, i) for i, j in enumerate('abcdefghijklmnopqrstuvwxyz'))
+        b32.update(dict((j, i) for i, j in enumerate('234567', 26)))
+
+        # some package hashes
+        tests = [
+            '35orsd4cenv743hg4i5vxha2lzayycby',
+            '6kfqtj7dap3773rxog6kkmoweix5gpwo',
+            'e6h6ff3uvmjbq3azik2ckr6ckwm3depv',
+            'snz2juf4ij7sv77cq3vs467q6acftmur',
+            '4eg47oedi5bbkhpoxw26v3oe6vamkfd7',
+            'vrwabwj6umeb5vjw6flx2rnft3j457rw']
+
+        for test_hash in tests:
+            # string containing raw bits of hash ('1' and '0')
+            expected = ''.join([format(b32[c], '#07b').replace('0b', '')
+                                for c in test_hash])
+
+            for bits in (1, 2, 3, 4, 7, 8, 9, 16, 64, 117, 128, 160):
+                actual_int = spack.spec.base32_prefix_bits(test_hash, bits)
+                fmt = "#0%sb" % (bits + 2)
+                actual = format(actual_int, fmt).replace('0b', '')
+
+                self.assertEqual(expected[:bits], actual)
+
+            self.assertRaises(
+                ValueError, spack.spec.base32_prefix_bits, test_hash, 161)
+
+            self.assertRaises(
+                ValueError, spack.spec.base32_prefix_bits, test_hash, 256)