From e5db9bb9b946d079216244975890cc008a3806a1 Mon Sep 17 00:00:00 2001
From: Michael Chirico <michaelchirico4@gmail.com>
Date: Thu, 18 Jul 2019 04:58:52 +0800
Subject: [PATCH] adds support for complex vectors to shift, part of #3690
 (#3693)

---
 NEWS.md               |  2 ++
 inst/tests/tests.Rraw |  9 ++++++++-
 src/shift.c           | 19 +++++++++++++++++++
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/NEWS.md b/NEWS.md
index d278665ad..fe20b7b1c 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -128,6 +128,8 @@
     identical(y1,y2) && identical(y1,y3)
     # TRUE
     ```
+    
+19. Extended support of `shift` to complex vectors, part of [#3690](https://github.com/Rdatatable/data.table/issues/3690).
 
 #### BUG FIXES
 
diff --git a/inst/tests/tests.Rraw b/inst/tests/tests.Rraw
index 6fb289dd5..469c0a59d 100644
--- a/inst/tests/tests.Rraw
+++ b/inst/tests/tests.Rraw
@@ -13358,7 +13358,7 @@ test(1963.07, shift(DT, -1:1),
           c(NA, 10L, 9L, 8L, 7L, 6L, 5L, 4L, 3L, 2L)))
 ## some coverage tests for good measure
 test(1963.08, shift(DT$x, type = 'some_other_type'), error='should be one of.*lag.*lead')
-test(1963.09, shift(c(1+3i, 2-1i)), error = 'Unsupported type')
+test(1963.09, shift(as.raw(0:1)), error = 'Unsupported type')
 test(1963.10, shift(DT, -1:1, type="shift", give.names = TRUE), # new type="shift" #3223
      ans <- list(`x_shift_-1` = c(2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L, 10L, NA),
           x_shift_0 = 1:10,
@@ -15314,6 +15314,13 @@ test(2066.09, DT[2:4, .(sum(B), mean(B)), by=A%%2L], data.table(A=0:1, V1=c(NA,
 test(2066.10, DT[,    .(sum(B, na.rm=TRUE), mean(B, na.rm=TRUE)), by=A%%2L], data.table(A=1:0, V1=c(-9i, -3i), V2=-3i))
 test(2066.11, DT[2:4, .(sum(B, na.rm=TRUE), mean(B, na.rm=TRUE)), by=A%%2L], data.table(A=0:1, V1=c(-3i, -3i), V2=-3i))
 
+# Shift complex values, part of #3690
+z = c(1:3) + c(3:1)*1i
+test(2067.1, shift(z), c(NA, z[1:2]))
+test(2067.2, shift(z, type = 'lead'), c(z[2:3], NA))
+test(2067.3, shift(z, fill = 1i), c(1i, z[1:2]))
+test(2067.4, shift(list(z, 1:3)), list(c(NA, z[1:2]), c(NA, 1:2)))
+
 
 ###################################
 #  Add new tests above this line  #
diff --git a/src/shift.c b/src/shift.c
index 10e7cbeb3..75eff7b2a 100644
--- a/src/shift.c
+++ b/src/shift.c
@@ -86,6 +86,25 @@ SEXP shift(SEXP obj, SEXP k, SEXP fill, SEXP type) {
       }
       break;
 
+    case CPLXSXP :
+      thisfill = PROTECT(coerceVector(fill, CPLXSXP)); protecti++;
+      Rcomplex cfill = COMPLEX(thisfill)[0];
+      for (int j=0; j<nk; j++) {
+        R_xlen_t thisk = MIN(abs(kd[j]), xrows);
+        SET_VECTOR_ELT(ans, i*nk+j, tmp=allocVector(CPLXSXP, xrows) );
+        Rcomplex *ctmp = COMPLEX(tmp);
+        size_t tailk = xrows-thisk;
+        if ((stype == LAG && kd[j] >= 0) || (stype == LEAD && kd[j] < 0)) {
+          if (tailk > 0) memmove(ctmp+thisk, COMPLEX(elem), tailk*size);
+          for (int m=0; m<thisk; m++) ctmp[m] = cfill;
+        } else {
+          if (tailk > 0) memmove(ctmp, COMPLEX(elem)+thisk, tailk*size);
+          for (int m=tailk; m<xrows; m++) ctmp[m] = cfill;
+        }
+        copyMostAttrib(elem, tmp);
+      }
+      break;
+
     case LGLSXP :
       thisfill = PROTECT(coerceVector(fill, LGLSXP)); protecti++;
       int lfill = LOGICAL(thisfill)[0];