Fix BoundedSequenceEncoding for big endian

viswanathgs · Aug 29, 2020 · cb92217 · cb92217
1 parent 3cfb3f6
commit cb92217
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/lm/interpolate/bounded_sequence_encoding.hh b/lm/interpolate/bounded_sequence_encoding.hh
@@ -10,10 +10,6 @@
 #include "util/exception.hh"
 #include "util/fixed_array.hh"
 
-#if BYTE_ORDER != LITTLE_ENDIAN
-#warning The interpolation code assumes little endian for now.
-#endif
-
 #include <algorithm>
 #include <cstring>
 
@@ -40,19 +36,28 @@ class BoundedSequenceEncoding {
         }
         cur |= static_cast<uint64_t>(*from) << i->shift;
       }
+#if BYTE_ORDER == BIG_ENDIAN
+      cur <<= (8 - overhang_) * 8;
+#endif
       memcpy(to, &cur, overhang_);
     }
 
     void Decode(const void *from_void, unsigned char *to) const {
       const uint8_t *from = static_cast<const uint8_t*>(from_void);
       uint64_t cur = 0;
       memcpy(&cur, from, first_copy_);
+#if BYTE_ORDER == BIG_ENDIAN
+      cur >>= (8 - first_copy_) * 8;
+#endif
       for (const Entry *i = entries_.begin(); i != entries_.end(); ++i, ++to) {
         if (UTIL_UNLIKELY(i->next)) {
           from += sizeof(uint64_t);
           cur = 0;
           std::memcpy(&cur, from,
               std::min<std::size_t>(sizeof(uint64_t), static_cast<const uint8_t*>(from_void) + byte_length_ - from));
+#if BYTE_ORDER == BIG_ENDIAN
+          cur >>= (8 - (static_cast<const uint8_t*>(from_void) + byte_length_ - from)) * 8;
+#endif
         }
         *to = (cur >> i->shift) & i->mask;
       }

diff --git a/lm/interpolate/bounded_sequence_encoding_test.cc b/lm/interpolate/bounded_sequence_encoding_test.cc
@@ -9,6 +9,17 @@ namespace lm {
 namespace interpolate {
 namespace {
 
+BOOST_AUTO_TEST_CASE(Simple) {
+  unsigned char bounds[] = {2};
+  BoundedSequenceEncoding enc(bounds, bounds + 1);
+  util::scoped_malloc backing(util::MallocOrThrow(enc.EncodedLength()));
+  unsigned char input = 1;
+  enc.Encode(&input, backing.get());
+  unsigned char output;
+  enc.Decode(backing.get(), &output);
+  BOOST_CHECK_EQUAL(1, output);
+}
+
 void ExhaustiveTest(unsigned char *bound_begin, unsigned char *bound_end) {
   BoundedSequenceEncoding enc(bound_begin, bound_end);
   util::scoped_malloc backing(util::MallocOrThrow(enc.EncodedLength()));