From 371ec2036fa55d667612c07f46c8b391ed0cbad4 Mon Sep 17 00:00:00 2001 From: Axel Kohlmeyer Date: Thu, 23 May 2024 01:24:41 -0400 Subject: [PATCH] support trailing brackets for sort() and rsort() --- src/variable.cpp | 30 ++++++++++++++++++++-------- src/variable.h | 3 ++- unittest/commands/test_variables.cpp | 7 +++++++ 3 files changed, 31 insertions(+), 9 deletions(-) diff --git a/src/variable.cpp b/src/variable.cpp index f499d7e6b1..e2c77fc520 100644 --- a/src/variable.cpp +++ b/src/variable.cpp @@ -2288,7 +2288,7 @@ double Variable::evaluate(char *str, Tree **tree, int ivar) if (math_function(word,contents,tree,treestack,ntreestack,argstack,nargstack,ivar)); else if (group_function(word,contents,tree,treestack,ntreestack,argstack,nargstack,ivar)); - else if (special_function(std::string(word),contents,tree,treestack,ntreestack,argstack,nargstack,ivar)); + else if (special_function(std::string(word),contents,tree,treestack,ntreestack,argstack,nargstack,ivar,str,i,ptr)); else if (feature_function(word,contents,tree,treestack,ntreestack,argstack,nargstack,ivar)); else print_var_error(FLERR,fmt::format("Invalid math/group/special/feature function '{}()' " "in variable formula", word),ivar); @@ -4271,7 +4271,7 @@ Region *Variable::region_function(char *id, int ivar) ------------------------------------------------------------------------- */ int Variable::special_function(const std::string &word, char *contents, Tree **tree, Tree **treestack, - int &ntreestack, double *argstack, int &nargstack, int ivar) + int &ntreestack, double *argstack, int &nargstack, int ivar, char *str, int &istr, char *&ptr) { double sx,sxx; double value,sy,sxy; @@ -4569,12 +4569,12 @@ int Variable::special_function(const std::string &word, char *contents, Tree **t if (method == SORT) std::sort(unsorted.begin(), unsorted.end(), std::less()); if (method == RSORT) std::sort(unsorted.begin(), unsorted.end(), std::greater()); - double *newvec; - memory->create(newvec,nvec,"variable:values"); - for (int m = 0; m < nvec; m++) - newvec[m] = unsorted[m]; - if (tree) { + double *newvec; + memory->create(newvec,nvec,"variable:values"); + for (int m = 0; m < nvec; m++) + newvec[m] = unsorted[m]; + auto newtree = new Tree(); newtree->type = VECTORARRAY; newtree->array = newvec; @@ -4582,8 +4582,22 @@ int Variable::special_function(const std::string &word, char *contents, Tree **t newtree->nstride = 1; newtree->selfalloc = 1; treestack[ntreestack++] = newtree; + } else { - error->all(FLERR, "Cannot use argstack for sorted vector"); + + // process one pair of trailing brackets + // point istr beyond last bracket + // nbracket = # of bracket pairs + // index = int inside each bracket pair, vector index + + if (str[istr] == '[') { + ptr = &str[istr]; + int index = int_between_brackets(ptr,1); + if (index > nvec) + print_var_error(FLERR, fmt::format("index {} exceeds vector size of {}\n", index, nvec),ivar); + istr = ptr-str+1; + argstack[nargstack++] = unsorted[index-1]; + } } } else { diff --git a/src/variable.h b/src/variable.h index fcc66461f6..f4a5e511c0 100644 --- a/src/variable.h +++ b/src/variable.h @@ -143,7 +143,8 @@ class Variable : protected Pointers { int math_function(char *, char *, Tree **, Tree **, int &, double *, int &, int); int group_function(char *, char *, Tree **, Tree **, int &, double *, int &, int); Region *region_function(char *, int); - int special_function(const std::string &, char *, Tree **, Tree **, int &, double *, int &, int); + int special_function(const std::string &, char *, Tree **, Tree **, int &, double *, int &, + int, char *, int &, char *&); int feature_function(char *, char *, Tree **, Tree **, int &, double *, int &, int); void peratom2global(int, char *, double *, int, tagint, Tree **, Tree **, int &, double *, int &); void custom2global(int *, double *, int, tagint, Tree **, Tree **, int &, double *, int &); diff --git a/unittest/commands/test_variables.cpp b/unittest/commands/test_variables.cpp index 4a0abd1856..1aa13ecf49 100644 --- a/unittest/commands/test_variables.cpp +++ b/unittest/commands/test_variables.cpp @@ -340,6 +340,9 @@ TEST_F(VariableTest, Expressions) command("variable vec4 vector '[1, 5, 2.5, -10, -5, 20, 120, 4, 3, 3]'"); command("variable sort vector sort(v_vec4)"); command("variable rsrt vector rsort(v_vec4)"); + command("variable max2 equal sort(v_vec4)[2]"); + command("variable rmax equal rsort(v_vec4)[1]"); + command("variable xxxl equal rsort(v_vec4)[11]"); command("variable isrt vector sort(v_one)"); variable->set("dummy index 1 2"); END_HIDE_OUTPUT(); @@ -372,6 +375,8 @@ TEST_F(VariableTest, Expressions) ASSERT_DOUBLE_EQ(variable->compute_equal("v_vec3"), 0.5); EXPECT_THAT(variable->retrieve("sort"), StrEq("[-10,-5,1,2.5,3,3,4,5,20,120]")); EXPECT_THAT(variable->retrieve("rsrt"), StrEq("[120,20,5,4,3,3,2.5,1,-5,-10]")); + ASSERT_DOUBLE_EQ(variable->compute_equal("v_max2"), -5); + ASSERT_DOUBLE_EQ(variable->compute_equal("v_rmax"), 120); TEST_FAILURE(".*ERROR: Variable six: Invalid thermo keyword 'XXX' in variable formula.*", command("print \"${six}\"");); @@ -385,6 +390,8 @@ TEST_F(VariableTest, Expressions) command("print \"${err3}\"");); TEST_FAILURE(".*ERROR: Variable one: Mis-matched special function variable in variable formula.*", command("print \"${isrt}\"");); + TEST_FAILURE(".*ERROR: Variable vec4: index 11 exceeds vector size of 10.*", + command("print \"${xxxl}\"");); } TEST_F(VariableTest, Functions)