2006-11-06 David Smith <dsmith@redhat.com>

	Added "limit EXP" support to foreach statements.
	* translate.cxx (c_tmpcounter::visit_foreach_loop): Handles
	"limit" member variable.
	(c_unparser::visit_foreach_loop): Ditto.
	* staptree.cxx (foreach_loop::print): Prints "limit EXP"
	addition.
	(traversing_visitor::visit_foreach_loop): Handles "limit" member
	variable.
	(deep_copy_visitor::visit_foreach_loop): Ditto.
	* staptree.h (struct foreach_loop): Added "limit" member
	variable.
	* stap.1.in: Added documentation for the "limit EXP" addition to
	foreach statement.
	* parse.cxx (lexer::scan): Added "limit" keyword for foreach
	statements.
	(parser::parse_foreach_loop): Parses "limit" keyword for foreach
	statements.
	* elaborate.cxx (symresolution_info::visit_foreach_loop): Handles
	"limit" member variable.
	(typeresolution_info::visit_foreach_loop): Ditto.

7 files changed, 157 insertions, 11 deletions

diff --git a/ChangeLog b/ChangeLog
index ccfaa942..4f6e56d7 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,26 @@
+2006-11-06  David Smith  <dsmith@redhat.com>
+
+	Added "limit EXP" support to foreach statements.
+	* translate.cxx (c_tmpcounter::visit_foreach_loop): Handles
+	"limit" member variable.
+	(c_unparser::visit_foreach_loop): Ditto.
+	* staptree.cxx (foreach_loop::print): Prints "limit EXP"
+	addition.
+	(traversing_visitor::visit_foreach_loop): Handles "limit" member
+	variable.
+	(deep_copy_visitor::visit_foreach_loop): Ditto.
+	* staptree.h (struct foreach_loop): Added "limit" member
+	variable.
+	* stap.1.in: Added documentation for the "limit EXP" addition to
+	foreach statement.
+	* parse.cxx (lexer::scan): Added "limit" keyword for foreach
+	statements.
+	(parser::parse_foreach_loop): Parses "limit" keyword for foreach
+	statements.
+	* elaborate.cxx (symresolution_info::visit_foreach_loop): Handles
+	"limit" member variable.
+	(typeresolution_info::visit_foreach_loop): Ditto.
+
 2006-11-03  Frank Ch. Eigler  <fche@elastic.org>
 
 	* tapsets.cxx (emit_address): Emit calls to _stp_module_relocate
diff --git a/elaborate.cxx b/elaborate.cxx
index f6c9bf9a..c532d0b0 100644
--- a/elaborate.cxx
+++ b/elaborate.cxx
@@ -1040,6 +1040,9 @@ symresolution_info::visit_foreach_loop (foreach_loop* e)
       hist->visit (this);
     }
 
+  if (e->limit)
+    e->limit->visit (this);
+
   e->block->visit (this);
 }
 
@@ -2375,6 +2378,12 @@ typeresolution_info::visit_foreach_loop (foreach_loop* e)
 	}
     }
 
+  if (e->limit)
+    {
+      t = pe_long;
+      e->limit->visit (this);
+    }
+
   t = pe_unknown;
   e->block->visit (this);  
 }
diff --git a/parse.cxx b/parse.cxx
index 2d31e334..dd41a259 100644
--- a/parse.cxx
+++ b/parse.cxx
@@ -543,6 +543,7 @@ lexer::scan ()
 	      || n->content == "for"
 	      || n->content == "foreach"
 	      || n->content == "in"
+	      || n->content == "limit"
 	      || n->content == "return"
 	      || n->content == "delete"
 	      || n->content == "while"
@@ -1518,6 +1519,7 @@ parser::parse_foreach_loop ()
   foreach_loop* s = new foreach_loop;
   s->tok = t;
   s->sort_direction = 0;
+  s->limit = NULL;
 
   t = next ();
   if (! (t->type == tok_operator && t->content == "("))
@@ -1591,6 +1593,13 @@ parser::parse_foreach_loop ()
       next();
     }
 
+  t = peek ();
+  if (tok_is(t, tok_keyword, "limit"))
+    {
+      next ();				// get past the "limit"
+      s->limit = parse_expression ();
+    }
+
   t = next ();
   if (! (t->type == tok_operator && t->content == ")"))
     throw parse_error ("expected ')'");
diff --git a/stap.1.in b/stap.1.in
index 7b572943..8b3f7409 100644
--- a/stap.1.in
+++ b/stap.1.in
@@ -304,15 +304,20 @@ While integer-valued EXP evaluates to non-zero, execute STMT.
 Execute EXP1 as initialization.  While EXP2 is non-zero, execute
 STMT, then the iteration expression EXP3.
 .TP
-.BR foreach " (VAR " in " ARRAY) STMT"
+.BR foreach " (VAR " in " ARRAY [ "limit " EXP ]) STMT"
 Loop over each element of the named global array, assigning current
 key to VAR.  The array may not be modified within the statement.
 By adding a single 
 .BR + " or " \-
 operator after the VAR or the ARRAY identifier, the iteration will
 proceed in a sorted order, by ascending or descending index or value.
+
+Using the optional
+.BR limit 
+keyword limits the number of loop iterations to EXP times.  EXP is
+evaluted once at the beginning of the loop.
 .TP
-.BR foreach " ([VAR1, VAR2, ...] " in " ARRAY) STMT"
+.BR foreach " ([VAR1, VAR2, ...] " in " ARRAY [ "limit " EXP ]) STMT"
 Same as above, used when the array is indexed with a tuple of keys.
 A sorting suffix may be used on at most one VAR or ARRAY identifier.
 .TP
diff --git a/staptree.cxx b/staptree.cxx
index 36712c12..fdfc516c 100644
--- a/staptree.cxx
+++ b/staptree.cxx
@@ -752,6 +752,11 @@ void foreach_loop::print (ostream& o) const
   base->print_indexable (o);
   if (sort_direction != 0 && sort_column == 0)
     o << (sort_direction > 0 ? "+" : "-");
+  if (limit)
+    {
+      o << " limit ";
+      limit->print (o);
+    }
   o << ") ";
   block->print (o);
 }
@@ -1320,6 +1325,10 @@ traversing_visitor::visit_foreach_loop (foreach_loop* s)
 
   for (unsigned i=0; i<s->indexes.size(); i++)
     s->indexes[i]->visit (this);
+
+  if (s->limit)
+    s->limit->visit (this);
+
   s->block->visit (this);
 }
 
@@ -1965,6 +1974,7 @@ deep_copy_visitor::visit_foreach_loop (foreach_loop* s)
 
   n->sort_direction = s->sort_direction;
   n->sort_column = s->sort_column;
+  require <expression*> (this, &(n->limit), s->limit);
 
   require <statement*> (this, &(n->block), s->block);
   provide <foreach_loop*> (this, n);
diff --git a/staptree.h b/staptree.h
index 0510bce3..413c65f1 100644
--- a/staptree.h
+++ b/staptree.h
@@ -464,6 +464,7 @@ struct foreach_loop: public statement
   indexable *base;
   int sort_direction; // -1: decreasing, 0: none, 1: increasing
   unsigned sort_column; // 0: value, 1..N: index
+  expression* limit; // optional iteration limit
 
   statement* block;
   void print (std::ostream& o) const;
diff --git a/translate.cxx b/translate.cxx
index 1f2e47f0..14d58ef2 100644
--- a/translate.cxx
+++ b/translate.cxx
@@ -2196,6 +2196,19 @@ c_tmpcounter::visit_foreach_loop (foreach_loop *s)
 	}
     }
 
+  // Create a temporary for the loop limit counter and the limit
+  // expression result.
+  if (s->limit)
+    {
+      tmpvar res_limit = parent->gensym (pe_long);
+      res_limit.declare(*parent);
+
+      s->limit->visit (this);
+
+      tmpvar limitv = parent->gensym (pe_long);
+      limitv.declare(*parent);
+    }
+
   s->block->visit (this);
 }
 
@@ -2223,6 +2236,14 @@ c_unparser::visit_foreach_loop (foreach_loop *s)
       
       // initialization
 
+      tmpvar *res_limit = NULL;
+      if (s->limit)
+        {
+	  // Evaluate the limit expression once.
+	  res_limit = new tmpvar(gensym(pe_long));
+	  c_assign (res_limit->qname(), s->limit, "foreach limit");
+	}
+      
       // aggregate array if required
       if (mv.is_parallel())
 	{
@@ -2231,20 +2252,43 @@ c_unparser::visit_foreach_loop (foreach_loop *s)
 	  o->indent(-1);
 
 	  // sort array if desired
-	  if (s->sort_direction) {
-	    o->newline() << "else"; // only sort if aggregation was ok
-	    o->newline(1) << "_stp_map_sort (" << mv.fetch_existing_aggregate() << ", "
-			  << s->sort_column << ", " << - s->sort_direction << ");";
-	    o->indent(-1);
-	  }
-	}
+	  if (s->sort_direction)
+	    {
+	      o->newline() << "else"; // only sort if aggregation was ok
+	      if (s->limit)
+	        {
+		  o->newline(1) << "_stp_map_sortn ("
+				<< mv.fetch_existing_aggregate() << ", "
+				<< *res_limit << ", " << s->sort_column << ", "
+				<< - s->sort_direction << ");";
+		}
+	      else
+	        {
+		  o->newline(1) << "_stp_map_sort ("
+				<< mv.fetch_existing_aggregate() << ", "
+				<< s->sort_column << ", "
+				<< - s->sort_direction << ");";
+		}
+	      o->indent(-1);
+	    }
+        }
       else
 	{      
 	  // sort array if desired
 	  if (s->sort_direction)
 	    {
-	      o->newline() << "_stp_map_sort (" << mv.qname() << ", "
-			   << s->sort_column << ", " << - s->sort_direction << ");";
+	      if (s->limit)
+	        {
+		  o->newline() << "_stp_map_sortn (" << mv.qname() << ", "
+			       << *res_limit << ", " << s->sort_column << ", "
+			       << - s->sort_direction << ");";
+		}
+	      else
+	        {
+		  o->newline() << "_stp_map_sort (" << mv.qname() << ", "
+			       << s->sort_column << ", "
+			       << - s->sort_direction << ");";
+		}
 	    }
 	}
 
@@ -2254,6 +2298,14 @@ c_unparser::visit_foreach_loop (foreach_loop *s)
 	aggregations_active.insert(mv.qname());
       o->newline() << iv << " = " << iv.start (mv) << ";";
       
+      tmpvar *limitv = NULL;
+      if (s->limit)
+      {
+	  // Create the loop limit variable here and initialize it.
+	  limitv = new tmpvar(gensym (pe_long));
+	  o->newline() << *limitv << " = 0LL;";
+      }
+
       // condition
       o->newline(-1) << toplabel << ":";
 
@@ -2270,6 +2322,18 @@ c_unparser::visit_foreach_loop (foreach_loop *s)
       loop_continue_labels.push_back (contlabel);
       o->newline() << "{";
       o->indent (1);
+
+      if (s->limit)
+      {
+	  // If we've been through LIMIT loop iterations, quit.
+	  o->newline() << "if (" << *limitv << "++ >= " << *res_limit
+		       << ") goto " << breaklabel << ";";
+
+	  // We're done with limitv and res_limit.
+	  delete limitv;
+	  delete res_limit;
+      }
+
       for (unsigned i = 0; i < s->indexes.size(); ++i)
 	{
 	  // copy the iter values into the specified locals
@@ -2307,11 +2371,36 @@ c_unparser::visit_foreach_loop (foreach_loop *s)
       var v = getvar(sym->referent, sym->tok);
       v.assert_hist_compatible(*hist);
 
+      tmpvar *res_limit = NULL;
+      tmpvar *limitv = NULL;
+      if (s->limit)
+        {
+	  // Evaluate the limit expression once.
+	  res_limit = new tmpvar(gensym(pe_long));
+	  c_assign (res_limit->qname(), s->limit, "foreach limit");
+
+	  // Create the loop limit variable here and initialize it.
+	  limitv = new tmpvar(gensym (pe_long));
+	  o->newline() << *limitv << " = 0LL;";
+	}
+      
       // XXX: break / continue don't work here yet
       o->newline() << "for (" << bucketvar << " = 0; " 
 		   << bucketvar << " < " << v.buckets() << "; "
 		   << bucketvar << "++) { ";
       o->newline(1);
+
+      if (s->limit)
+      {
+	  // If we've been through LIMIT loop iterations, quit.
+	  o->newline() << "if (" << *limitv << "++ >= " << *res_limit
+		       << ") break;";
+
+	  // We're done with limitv and res_limit.
+	  delete limitv;
+	  delete res_limit;
+      }
+
       s->block->visit (this);
       o->newline(-1) << "}";
     }