assign: factor out index fixup

Instead of walking the attribute list directly, use R_mapAttrib().
Create a hash table of index names instead of relying on chin() and a
temporary string vector. Move all temporary allocations onto the R heap.

Author: aitap

src/assign.c

@@ -256,6 +256,99 @@ SEXP selfrefokwrapper(SEXP x, SEXP verbose) {
   return ScalarInteger(_selfrefok(x,FALSE,LOGICAL(verbose)[0]));
 }
 
+struct attrib_name_ctx {
+  hashtab *indexNames;
+  SEXP index, assignedNames;
+  R_xlen_t len;
+  bool verbose;
+};
+static SEXP getOneAttribName(SEXP key, SEXP val, void *ctx_) {
+  (void)val;
+  struct attrib_name_ctx *ctx = ctx_;
+  if (ctx->indexNames)
+    hash_set(ctx->indexNames, PRINTNAME(key), 1);
+  else
+    ctx->len++;
+  return NULL;
+}
+
+static SEXP fixIndexAttrib(SEXP tag, SEXP value, void *ctx_) {
+  const struct attrib_name_ctx *ctx = ctx_;
+
+  hashtab *indexNames = ctx->indexNames;
+  SEXP index = ctx->index, assignedNames = ctx->assignedNames;
+  R_xlen_t indexLength = xlength(value);
+  bool verbose = ctx->verbose;
+
+  const char *tc1, *c1;
+  tc1 = c1 = CHAR(PRINTNAME(tag));  // the index name; e.g. "__col1__col2"
+
+  if (*tc1!='_' || *(tc1+1)!='_') {
+    // fix for #1396
+    if (verbose) {
+      Rprintf(_("Dropping index '%s' as it doesn't have '__' at the beginning of its name. It was very likely created by v1.9.4 of data.table.\n"), tc1);
+    }
+    setAttrib(index, tag, R_NilValue);
+    return NULL;
+  }
+
+  tc1 += 2; // tc1 always marks the start of a key column
+  if (!*tc1) internal_error(__func__, "index name ends with trailing __"); // # nocov
+
+  void *vmax = vmaxget();
+  // check the position of the first appearance of an assigned column in the index.
+  // the new index will be truncated to this position.
+  size_t newKeyLength = strlen(c1);
+  char *s4 = R_alloc(newKeyLength + 3, 1);
+  memcpy(s4, c1, newKeyLength);
+  memcpy(s4 + newKeyLength, "__", 3);
+
+  for(int i = 0; i < xlength(assignedNames); i++){
+    const char *tc2 = CHAR(STRING_ELT(assignedNames, i));
+    void *vmax2 = vmaxget();
+    size_t tc2_len = strlen(tc2);
+    char *s5 = R_alloc(tc2_len + 5, 1); //4 * '_' + \0
+    memcpy(s5, "__", 2);
+    memcpy(s5 + 2, tc2, tc2_len);
+    memcpy(s5 + 2 + tc2_len, "__", 3);
+    tc2 = strstr(s4, s5);
+    if(tc2 && (tc2 - s4 < newKeyLength)){ // new column is part of key; match is before last match
+      newKeyLength = tc2 - s4;
+    }
+    vmaxset(vmax2);
+  }
+
+  s4[newKeyLength] = '\0'; // truncate the new key to the new length
+  if(newKeyLength == 0){ // no valid key column remains. Drop the key
+    setAttrib(index, tag, R_NilValue);
+    hash_set(indexNames, PRINTNAME(tag), 0);
+    if (verbose) {
+      Rprintf(_("Dropping index '%s' due to an update on a key column\n"), c1+2);
+    }
+  } else if(newKeyLength < strlen(c1)) {
+    SEXP s4Str = PROTECT(mkChar(s4));
+    if(indexLength == 0 && // shortened index can be kept since it is just information on the order (see #2372)
+      !hash_lookup(indexNames, s4Str, 0)) { // index with shortened name not present yet
+      setAttrib(index, installChar(s4Str), value);
+      hash_set(indexNames, PRINTNAME(tag), 0);
+      setAttrib(index, tag, R_NilValue);
+      hash_set(indexNames, s4Str, 1);
+      if (verbose)
+        Rprintf(_("Shortening index '%s' to '%s' due to an update on a key column\n"), c1+2, s4+2);
+    } else { // indexLength > 0 || shortened name present already
+      // indexLength > 0 indicates reordering. Drop it to avoid spurious reordering in non-indexed columns (#2372)
+      // shortened name already present indicates that index needs to be dropped to avoid duplicate indices.
+      setAttrib(index, tag, R_NilValue);
+      hash_set(indexNames, tag, 0);
+      if (verbose)
+        Rprintf(_("Dropping index '%s' due to an update on a key column\n"), c1+2);
+    }
+    UNPROTECT(1); // s4Str
+  } //else: index is not affected by assign: nothing to be done
+  vmaxset(vmax);
+  return NULL;
+}
+
 int *_Last_updated = NULL;
 
 SEXP assign(SEXP dt, SEXP rows, SEXP cols, SEXP newcolnames, SEXP values)
@@ -265,11 +358,11 @@ SEXP assign(SEXP dt, SEXP rows, SEXP cols, SEXP newcolnames, SEXP values)
   // cols : column names or numbers corresponding to the values to set
   // rows : row numbers to assign
   R_len_t numToDo, targetlen, vlen, oldncol, oldtncol, coln, protecti=0, newcolnum, indexLength;
-  SEXP targetcol, nullint, s, colnam, tmp, key, index, a, assignedNames, indexNames;
+  SEXP targetcol, nullint, s, colnam, tmp, key, index, a, assignedNames;
   bool verbose=GetVerbose();
   int ndelete=0;  // how many columns are being deleted
   const char *c1, *tc1, *tc2;
-  int *buf, indexNo;
+  int *buf;
   if (isNull(dt)) error(_("assign has been passed a NULL dt"));
   if (TYPEOF(dt) != VECSXP) error(_("dt passed to assign isn't type VECSXP"));
   if (islocked(dt))
@@ -549,93 +642,17 @@ SEXP assign(SEXP dt, SEXP rows, SEXP cols, SEXP newcolnames, SEXP values)
   }
   index = getAttrib(dt, install("index"));
   if (index != R_NilValue) {
-    s = ATTRIB(index);
-    indexNo = 0;
-    // get a vector with all index names
-    PROTECT(indexNames = allocVector(STRSXP, xlength(s))); protecti++;
-    while(s != R_NilValue){
-      SET_STRING_ELT(indexNames, indexNo, PRINTNAME(TAG(s)));
-      indexNo++;
-      s = CDR(s);
-    }
-    s = ATTRIB(index); // reset to first element
-    indexNo = 0;
-    while(s != R_NilValue) {
-      a = TAG(s);
-      indexLength = xlength(CAR(s));
-      tc1 = c1 = CHAR(PRINTNAME(a));  // the index name; e.g. "__col1__col2"
-      if (*tc1!='_' || *(tc1+1)!='_') {
-        // fix for #1396
-        if (verbose) {
-          Rprintf(_("Dropping index '%s' as it doesn't have '__' at the beginning of its name. It was very likely created by v1.9.4 of data.table.\n"), tc1);
-        }
-        setAttrib(index, a, R_NilValue);
-        indexNo++;
-        s = CDR(s);
-        continue; // with next index
-      }
-      tc1 += 2; // tc1 always marks the start of a key column
-      if (!*tc1) internal_error(__func__, "index name ends with trailing __"); // # nocov
-      // check the position of the first appearance of an assigned column in the index.
-      // the new index will be truncated to this position.
-      char *s4 = malloc(strlen(c1) + 3);
-      if (!s4) {
-        internal_error(__func__, "Couldn't allocate memory for s4"); // # nocov
-      }
-      memcpy(s4, c1, strlen(c1));
-      memset(s4 + strlen(c1), '\0', 1);
-      strcat(s4, "__"); // add trailing '__' to newKey so we can search for pattern '__colName__' also at the end of the index.
-      int newKeyLength = strlen(c1);
-      for(int i = 0; i < xlength(assignedNames); i++){
-        tc2 = CHAR(STRING_ELT(assignedNames, i));
-        char *s5 = malloc(strlen(tc2) + 5); //4 * '_' + \0
-        if (!s5) {
-          free(s4);                                                  // # nocov
-          internal_error(__func__, "Couldn't allocate memory for s5"); // # nocov
-        }
-        memset(s5, '_', 2);
-        memset(s5 + 2, '\0', 1);
-        strcat(s5, tc2);
-        strcat(s5, "__");
-        tc2 = strstr(s4, s5);
-        if(tc2 == NULL){ // column is not part of key
-          free(s5);
-          continue;
-        }
-        if(tc2 - s4 < newKeyLength){ // new column match is before last match
-          newKeyLength = tc2 - s4;
-        }
-        free(s5);
-      }
-      memset(s4 + newKeyLength, '\0', 1); // truncate the new key to the new length
-      if(newKeyLength == 0){ // no valid key column remains. Drop the key
-        setAttrib(index, a, R_NilValue);
-        SET_STRING_ELT(indexNames, indexNo, NA_STRING);
-        if (verbose) {
-          Rprintf(_("Dropping index '%s' due to an update on a key column\n"), c1+2);
-        }
-      } else if(newKeyLength < strlen(c1)) {
-        SEXP s4Str = PROTECT(mkString(s4));
-        if(indexLength == 0 && // shortened index can be kept since it is just information on the order (see #2372)
-           LOGICAL(chin(s4Str, indexNames))[0] == 0) {// index with shortened name not present yet
-          SET_TAG(s, install(s4));
-          SET_STRING_ELT(indexNames, indexNo, mkChar(s4));
-          if (verbose)
-            Rprintf(_("Shortening index '%s' to '%s' due to an update on a key column\n"), c1+2, s4 + 2);
-        } else { // indexLength > 0 || shortened name present already
-          // indexLength > 0 indicates reordering. Drop it to avoid spurious reordering in non-indexed columns (#2372)
-          // shortened name already present indicates that index needs to be dropped to avoid duplicate indices.
-          setAttrib(index, a, R_NilValue);
-          SET_STRING_ELT(indexNames, indexNo, NA_STRING);
-          if (verbose)
-            Rprintf(_("Dropping index '%s' due to an update on a key column\n"), c1+2);
-        }
-        UNPROTECT(1); // s4Str
-      } //else: index is not affected by assign: nothing to be done
-      free(s4);
-      indexNo ++;
-      s = CDR(s);
-    }
+    struct attrib_name_ctx ctx = { 0, };
+    R_mapAttrib(index, getOneAttribName, &ctx); // how many attributes?
+    hashtab *h = hash_create(ctx.len);
+    PROTECT(h->prot);
+    ctx.indexNames = h;
+    R_mapAttrib(index, getOneAttribName, &ctx); // now remember the names
+    ctx.index = index;
+    ctx.assignedNames = assignedNames;
+    ctx.verbose = verbose;
+    R_mapAttrib(index, fixIndexAttrib, &ctx); // adjust indices as needed
+    UNPROTECT(1); // h
   }
   if (ndelete) {
     // delete any columns assigned NULL (there was a 'continue' earlier in loop above)