avutil/eval: Add av_expr_count_func() similar to av_expr_count_vars()

Signed-off-by: Michael Niedermayer <michael@niedermayer.cc>

1 files changed, 20 insertions, 8 deletions

diff --git a/libavutil/eval.c b/libavutil/eval.c
index 62d2ae938b..d527f6a9d0 100644
--- a/libavutil/eval.c
+++ b/libavutil/eval.c
@@ -166,8 +166,8 @@ struct AVExpr {
         e_sgn,
     } type;
     double value; // is sign in other types
+    int const_index;
     union {
-        int const_index;
         double (*func0)(double);
         double (*func1)(void *, double);
         double (*func2)(void *, double, double);
@@ -185,7 +185,7 @@ static double eval_expr(Parser *p, AVExpr *e)
 {
     switch (e->type) {
         case e_value:  return e->value;
-        case e_const:  return e->value * p->const_values[e->a.const_index];
+        case e_const:  return e->value * p->const_values[e->const_index];
         case e_func0:  return e->value * e->a.func0(eval_expr(p, e->param[0]));
         case e_func1:  return e->value * e->a.func1(p->opaque, eval_expr(p, e->param[0]));
         case e_func2:  return e->value * e->a.func2(p->opaque, eval_expr(p, e->param[0]), eval_expr(p, e->param[1]));
@@ -367,7 +367,7 @@ static int parse_primary(AVExpr **e, Parser *p)
         if (strmatch(p->s, p->const_names[i])) {
             p->s+= strlen(p->const_names[i]);
             d->type = e_const;
-            d->a.const_index = i;
+            d->const_index = i;
             *e = d;
             return 0;
         }
@@ -478,6 +478,7 @@ static int parse_primary(AVExpr **e, Parser *p)
             if (strmatch(next, p->func1_names[i])) {
                 d->a.func1 = p->funcs1[i];
                 d->type = e_func1;
+                d->const_index = i;
                 *e = d;
                 return 0;
             }
@@ -487,6 +488,7 @@ static int parse_primary(AVExpr **e, Parser *p)
             if (strmatch(next, p->func2_names[i])) {
                 d->a.func2 = p->funcs2[i];
                 d->type = e_func2;
+                d->const_index = i;
                 *e = d;
                 return 0;
             }
@@ -735,22 +737,32 @@ end:
     return ret;
 }
 
-int av_expr_count_vars(AVExpr *e, unsigned *counter, int size)
+static int expr_count(AVExpr *e, unsigned *counter, int size, int type)
 {
     int i;
 
     if (!e || !counter || !size)
         return AVERROR(EINVAL);
 
-    for (i = 0; e->type != e_const && i < 3 && e->param[i]; i++)
-        av_expr_count_vars(e->param[i], counter, size);
+    for (i = 0; e->type != type && i < 3 && e->param[i]; i++)
+        expr_count(e->param[i], counter, size, type);
 
-    if (e->type == e_const && e->a.const_index < size)
-        counter[e->a.const_index]++;
+    if (e->type == type && e->const_index < size)
+        counter[e->const_index]++;
 
     return 0;
 }
 
+int av_expr_count_vars(AVExpr *e, unsigned *counter, int size)
+{
+    return expr_count(e, counter, size, e_const);
+}
+
+int av_expr_count_func(AVExpr *e, unsigned *counter, int size, int arg)
+{
+    return expr_count(e, counter, size, ((int[]){e_const, e_func1, e_func2})[arg]);
+}
+
 double av_expr_eval(AVExpr *e, const double *const_values, void *opaque)
 {
     Parser p = { 0 };