From c97f8906511dbb97079eaf44329db5aa3216470a Mon Sep 17 00:00:00 2001
From: Matthijs Douze <matthijs@fb.com>
Date: Fri, 28 Aug 2020 23:42:07 -0700
Subject: [PATCH] make sure swig_ptr and rev_swig_ptr work on all primitive
 types (#1382)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/1382

The array types supported for swig_ptr were not complete. This diff fixes that.

Reviewed By: wickedfoo

Differential Revision: D23411297

fbshipit-source-id: d94249b140aeb8c8179d9da3fcbc97eb034eac91
---
 faiss/python/swigfaiss.swig | 28 +++++++++++++++++++++++-----
 tests/test_build_blocks.py  | 20 ++++++++++++++++++++
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig
index 290319ba8..32107d848 100644
--- a/faiss/python/swigfaiss.swig
+++ b/faiss/python/swigfaiss.swig
@@ -210,6 +210,7 @@ namespace std {
 %include <std_pair.i>
 %include <std_map.i>
 
+// primitive array types
 %template(FloatVector) std::vector<float>;
 %template(DoubleVector) std::vector<double>;
 %template(ByteVector) std::vector<uint8_t>;
@@ -220,6 +221,10 @@ namespace std {
 %template(LongVector) std::vector<long>;
 %template(LongLongVector) std::vector<long long>;
 %template(IntVector) std::vector<int>;
+%template(UInt32Vector) std::vector<uint32_t>;
+%template(Int16Vector) std::vector<int16_t>;
+%template(UInt16Vector) std::vector<uint16_t>;
+
 %template(FloatVectorVector) std::vector<std::vector<float> >;
 %template(ByteVectorVector) std::vector<std::vector<unsigned char> >;
 %template(LongVectorVector) std::vector<std::vector<long> >;
@@ -799,15 +804,24 @@ PyObject *swig_ptr (PyObject *a)
     if(PyArray_TYPE(ao) == NPY_FLOAT64) {
         return SWIG_NewPointerObj(data, SWIGTYPE_p_double, 0);
     }
-    if(PyArray_TYPE(ao) == NPY_INT32) {
-        return SWIG_NewPointerObj(data, SWIGTYPE_p_int, 0);
-    }
     if(PyArray_TYPE(ao) == NPY_UINT8) {
         return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_char, 0);
     }
     if(PyArray_TYPE(ao) == NPY_INT8) {
         return SWIG_NewPointerObj(data, SWIGTYPE_p_char, 0);
     }
+    if(PyArray_TYPE(ao) == NPY_UINT16) {
+        return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_short, 0);
+    }
+    if(PyArray_TYPE(ao) == NPY_INT16) {
+        return SWIG_NewPointerObj(data, SWIGTYPE_p_short, 0);
+    }
+    if(PyArray_TYPE(ao) == NPY_UINT32) {
+        return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_int, 0);
+    }
+    if(PyArray_TYPE(ao) == NPY_INT32) {
+        return SWIG_NewPointerObj(data, SWIGTYPE_p_int, 0);
+    }
     if(PyArray_TYPE(ao) == NPY_UINT64) {
 #ifdef SWIGWORDSIZE64
         return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_long, 0);
@@ -870,8 +884,13 @@ PyObject * rev_swig_ptr(ctype *src, size_t size);
 %enddef
 
 REV_SWIG_PTR(float, NPY_FLOAT32);
-REV_SWIG_PTR(int, NPY_INT32);
+REV_SWIG_PTR(double, NPY_FLOAT64);
 REV_SWIG_PTR(unsigned char, NPY_UINT8);
+REV_SWIG_PTR(char, NPY_INT8);
+REV_SWIG_PTR(unsigned short, NPY_UINT16);
+REV_SWIG_PTR(short, NPY_INT16);
+REV_SWIG_PTR(int, NPY_INT32);
+REV_SWIG_PTR(unsigned int, NPY_UINT32);
 REV_SWIG_PTR(int64_t, NPY_INT64);
 REV_SWIG_PTR(uint64_t, NPY_UINT64);
 
@@ -1038,4 +1057,3 @@ struct MapLong2Long {
  %}
 
 // End of file...
- 
diff --git a/tests/test_build_blocks.py b/tests/test_build_blocks.py
index 2922f892f..09b4bc84f 100644
--- a/tests/test_build_blocks.py
+++ b/tests/test_build_blocks.py
@@ -628,5 +628,25 @@ class TestSWIGWrap(unittest.TestCase):
 
         assert isinstance(index2, faiss.IndexRefineFlat)
 
+    def do_test_array_type(self, dtype):
+        """ tests swig_ptr and rev_swig_ptr for this type of array """
+        a = np.arange(12).astype(dtype)
+        ptr = faiss.swig_ptr(a)
+        print(ptr)
+        a2 = faiss.rev_swig_ptr(ptr, 12)
+        np.testing.assert_array_equal(a, a2)
+
+    def test_all_array_types(self):
+        self.do_test_array_type('float32')
+        self.do_test_array_type('float64')
+        self.do_test_array_type('int8')
+        self.do_test_array_type('uint8')
+        self.do_test_array_type('int16')
+        self.do_test_array_type('uint16')
+        self.do_test_array_type('int32')
+        self.do_test_array_type('uint32')
+        self.do_test_array_type('int64')
+        self.do_test_array_type('uint64')
+
 if __name__ == '__main__':
     unittest.main()