From 00bc2b630dd27eae352b99bb69786516d2c3f37b Mon Sep 17 00:00:00 2001
From: KHLee <khlee870529@gmail.com>
Date: Wed, 22 Jan 2025 21:05:56 +0800
Subject: [PATCH] Make return array of argsort into uint64_t and wrap argsort

In order to make the datatype of return value of SimpleArray::argsort
compatible to the wrapped python datatypes, the return value is changed
to SimpleArray<uint64_t>.
Function argsort is also wrapped in this commit.
---
 cpp/modmesh/buffer/SimpleArray.hpp            | 12 ++++++------
 cpp/modmesh/buffer/pymod/wrap_SimpleArray.cpp |  4 ++++
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/cpp/modmesh/buffer/SimpleArray.hpp b/cpp/modmesh/buffer/SimpleArray.hpp
index 29fd701c..953e8345 100644
--- a/cpp/modmesh/buffer/SimpleArray.hpp
+++ b/cpp/modmesh/buffer/SimpleArray.hpp
@@ -595,28 +595,28 @@ class SimpleArray
         std::sort(begin(), end());
     }
 
-    SimpleArray<size_t> argsort(void)
+    SimpleArray<uint64_t> argsort(void)
     {
         if (ndim() != 1){
             throw std::runtime_error("SimpleArray: Sorting is only supported in 1D array.");
         }
 
-        SimpleArray<size_t> ret(shape());
+        SimpleArray<uint64_t> ret(shape());
 
         {   // Return array initialization
-            size_t cnt = 0;
-            std::for_each(ret.begin(), ret.end(), [&cnt](size_t &v){v = cnt++;});
+            uint64_t cnt = 0;
+            std::for_each(ret.begin(), ret.end(), [&cnt](uint64_t &v){v = cnt++;});
         }   
 
         value_type const *buf = body();
-        auto cmp = [buf](size_t a, size_t b) {
+        auto cmp = [buf](uint64_t a, uint64_t b) {
             return buf[a] < buf[b];
         };
         std::sort(ret.begin(), ret.end(), cmp);
         return ret;
     }
 
-    void apply_argsort(SimpleArray<size_t> const &sorted_args)
+    void apply_argsort(SimpleArray<uint64_t> const &sorted_args)
     {
         if (ndim() != 1 || sorted_args.ndim() != 1){
             throw std::runtime_error("SimpleArray: Sorting is only supported in 1D array.");
diff --git a/cpp/modmesh/buffer/pymod/wrap_SimpleArray.cpp b/cpp/modmesh/buffer/pymod/wrap_SimpleArray.cpp
index b5902e26..d0655a96 100644
--- a/cpp/modmesh/buffer/pymod/wrap_SimpleArray.cpp
+++ b/cpp/modmesh/buffer/pymod/wrap_SimpleArray.cpp
@@ -151,6 +151,10 @@ class MODMESH_PYTHON_WRAPPER_VISIBILITY WrapSimpleArray
                 "sort",
                 [](wrapped_type & self)
                 { self.sort(); })
+            .def(
+                "argsort",
+                [](wrapped_type & self)
+                { return pybind11::cast(self.argsort()); })
             .def_property_readonly("has_ghost", &wrapped_type::has_ghost)
             .def_property("nghost", &wrapped_type::nghost, &wrapped_type::set_nghost)
             .def_property_readonly("nbody", &wrapped_type::nbody)