Allow IntegerVector subsetting (0-based); add unit tests; use VectorBase so sugar exprs work

Author: kevinushey

inst/include/Rcpp/vector/Subsetter.h

@@ -72,20 +72,30 @@ class Subsetter {
   inline Vector<RTYPE, StoragePolicy> subset_impl( const VECTOR& this_, const Vector<LGLSXP, OtherStoragePolicy>& x ) const {
     Vector<INTSXP, StoragePolicy> tmp = which_na(x);
     if (!tmp.size()) return Vector<RTYPE, StoragePolicy>(0);
-    else return subset_impl__hidden(this_, tmp);
+    else return subset_impl(this_, tmp);
   }
 
   // Subsetting for characters
   template <template <class> class OtherStoragePolicy>
   inline Vector<RTYPE, StoragePolicy> subset_impl( const VECTOR& this_, const Vector<STRSXP, OtherStoragePolicy>& x ) const {
+      
+      if (Rf_isNull( Rf_getAttrib(this_, R_NamesSymbol) )) {
+          int n = x.size();
+          Vector<RTYPE, StoragePolicy> output(n);
+          for (int i=0; i < n; ++i) {
+              output[i] = traits::get_na<RTYPE>();
+          }
+          return output;
+      }
+      
     Vector<STRSXP, StoragePolicy> names = as< Vector<STRSXP, StoragePolicy> >(Rf_getAttrib(this_, R_NamesSymbol));
     Vector<INTSXP, StoragePolicy> idx = match(x, names); // match returns 1-based index
     // apparently, we don't see sugar, so we have to do this manually
     Vector<INTSXP, StoragePolicy> idxm1 = no_init(idx.size());
     for (int i=0; i < idx.size(); ++i) {
         idxm1[i] = idx[i] - 1;
     }
-    Vector<RTYPE, StoragePolicy> output = subset_impl__hidden(this_, idxm1);
+    Vector<RTYPE, StoragePolicy> output = subset_impl(this_, idxm1);
     int n = output.size();
     if (n == 0) return Vector<RTYPE, StoragePolicy>(0);
     Vector<STRSXP, StoragePolicy> out_names = no_init(n);
@@ -100,21 +110,11 @@ class Subsetter {
   const VECTOR& vec;
   const T& other;
 
-  private:
-  
-  // Subsetting for integers -- note that it is 0-based
-  template <template <class> class OtherStoragePolicy>
-  inline Vector<RTYPE, StoragePolicy> 
-  subset_impl( const VECTOR this_, const Vector<INTSXP, OtherStoragePolicy>& x ) const {
-    stop("Subset is not yet implemented for IntegerVectors");
-    return Vector<RTYPE, StoragePolicy>();
-  }
-  
   // Subsetting for integers -- note that it is 0-based
   // hidden until we decide whether to go with 0 or 1 based indexing
   template <template <class> class OtherStoragePolicy>
   inline Vector<RTYPE, StoragePolicy> 
-  subset_impl__hidden( const VECTOR this_, const Vector<INTSXP, OtherStoragePolicy>& x ) const {
+  subset_impl( const VECTOR this_, const Vector<INTSXP, OtherStoragePolicy>& x ) const {
     int n = x.size();
     if (n == 0) return this_;
     Vector<RTYPE, StoragePolicy> output = no_init(n);
@@ -124,8 +124,12 @@ class Subsetter {
       else if (x[i] > this_.size() - 1) output[i] = traits::get_na<RTYPE>();
       else output[i] = (this_)[ x[i] ];
     }
+    
     if (!Rf_isNull( Rf_getAttrib( this_, R_NamesSymbol) )) {
-      Vector<STRSXP, StoragePolicy> thisnames = as<Vector<STRSXP, StoragePolicy> >(Rf_getAttrib(this_, R_NamesSymbol));
+      
+      Vector<STRSXP, StoragePolicy> thisnames = 
+        as<Vector<STRSXP, StoragePolicy> >(Rf_getAttrib(this_, R_NamesSymbol));
+        
       Vector<STRSXP, StoragePolicy> outnames = no_init(n);
       for (int i=0; i < n; ++i) {
         if (x[i] == NA_INTEGER) outnames[i] = NA_STRING;

inst/include/Rcpp/vector/Vector.h

@@ -313,11 +313,12 @@ class Vector :
         return RObject( Storage::get__() );
     }
 
-    template <int RTYPE2>
-    Vector operator[](const Vector<RTYPE2>& rhs) const {
-        return Subsetter<RTYPE, StoragePolicy, Vector<RTYPE2, StoragePolicy> >(*this, rhs);
+    // sugar subsetting requires dispatch on VectorBase
+    template <int RTYPE2, bool na, typename T>
+    Vector operator[](const VectorBase<RTYPE2, na, T>& rhs) const {
+        return Subsetter<RTYPE, StoragePolicy, Vector<RTYPE2, PreserveStorage> >(*this, rhs);
     }
-	
+    
     Vector& sort(){
         typename traits::storage_type<RTYPE>::type* start = internal::r_vector_start<RTYPE>( Storage::get__() ) ;
         std::sort( 

inst/unitTests/cpp/Subset.cpp

@@ -0,0 +1,42 @@
+#include <Rcpp.h>
+using namespace Rcpp;
+
+// [[Rcpp::export]]
+NumericVector subset_test_int(NumericVector x, IntegerVector y) {
+    return x[y];
+}
+
+// [[Rcpp::export]]
+NumericVector subset_test_lgcl(NumericVector x, LogicalVector y) {
+    return x[y];
+}
+
+// [[Rcpp::export]]
+NumericVector subset_test_char(NumericVector x, CharacterVector y) {
+    return x[y];
+}
+
+// [[Rcpp::export]]
+List subset_test_list(List x, CharacterVector y) {
+    return x[y];
+}
+
+// [[Rcpp::export]]
+List subset_test_list_int(List x, IntegerVector y) {
+    return x[y];
+}
+
+// [[Rcpp::export]]
+List subset_test_list_lgcl(List x, LogicalVector y) {
+    return x[y];
+}
+
+// [[Rcpp::export]]
+NumericVector subset_test_greater_0(NumericVector x) {
+    return x[ x > 0 ];
+}
+
+// [[Rcpp::export]]
+List subset_test_literal(List x) {
+    return x["foo"];
+}

inst/unitTests/runit.subset.R

@@ -0,0 +1,66 @@
+#!/usr/bin/r -t
+# -*- mode: R; tab-width: 4; -*-
+#
+# Copyright (C) 2014  Dirk Eddelbuettel, Romain Francois and Kevin Ushey
+#
+# This file is part of Rcpp.
+#
+# Rcpp is free software: you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 2 of the License, or
+# (at your option) any later version.
+#
+# Rcpp is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with Rcpp.  If not, see <http://www.gnu.org/licenses/>.
+
+.runThisTest <- Sys.getenv("RunAllRcppTests") == "yes"
+
+if (.runThisTest) {
+
+    .setUp <- Rcpp:::unitTestSetup("Subset.cpp")
+    
+    test.subset <- function() {
+        
+        x <- rnorm(5)
+        names(x) <- letters[1:5]
+        
+        checkIdentical( x[c(1, 2, 10)], subset_test_int(x, c(0L, 1L, 9L)),
+            "integer subsetting")
+        
+        checkIdentical( x[ c('b', 'a') ], subset_test_char(x, c('b', 'a')),
+            "character subsetting")
+        
+        checkIdentical( c(1, 2, 3)['a'], subset_test_char( c(1, 2, 3), 'a' ),
+            "character subsetting -- no names on x")
+        
+        lgcl <- c(TRUE, FALSE, NA, TRUE, TRUE)
+        checkIdentical( 
+            x[lgcl],
+            subset_test_lgcl(x, lgcl),
+            "logical subsetting" )
+        
+        names(x) <- c('a', 'b', 'b', 'c', 'd')
+        checkIdentical(
+            x['b'],
+            subset_test_char(x, 'b'),
+            "character subsetting -- duplicated name")
+        
+        l <- as.list(x)
+        checkIdentical(
+            l[c('b', 'Q')],
+            subset_test_list(x, c('b', 'Q')),
+            "list subsetting")
+        
+        checkIdentical(
+            x[ x > 0 ],
+            subset_test_greater_0(x),
+            "sugar subsetting (x[x > 0])")
+        
+    }
+
+}