2 files changed, 29 insertions, 1 deletions
diff --git a/src/runtime/c/pgf/expr.c b/src/runtime/c/pgf/expr.c
index 8fee28fb9..c1f803385 100644
--- a/src/runtime/c/pgf/expr.c
+++ b/src/runtime/c/pgf/expr.c
@@ -1,11 +1,12 @@
 #include "pgf.h"
+#include "data.h"
 #include <gu/assert.h>
 #include <gu/utf8.h>
 #include <gu/seq.h>
 #include <ctype.h>
 #include <stdio.h>
 #include <stdlib.h>
-
+#include <math.h>
 
 static PgfExpr
 pgf_expr_unwrap(PgfExpr expr)
@@ -1500,3 +1501,27 @@ pgf_type_eq(PgfType* t1, PgfType* t2)
 
 	return true;
 }
+
+PGF_API prob_t
+pgf_compute_tree_probability(PgfPGF *gr, PgfExpr expr)
+{
+	GuVariantInfo ei = gu_variant_open(expr);
+	switch (ei.tag) {
+	case PGF_EXPR_APP: {
+		PgfExprApp* app = ei.data;
+		return pgf_compute_tree_probability(gr, app->fun) +
+		       pgf_compute_tree_probability(gr, app->arg);
+	}
+	case PGF_EXPR_FUN: {
+		PgfExprFun* fun = ei.data;
+		PgfAbsFun* absfun =
+			gu_seq_binsearch(gr->abstract.funs, pgf_absfun_order, PgfAbsFun, fun->fun);
+		if (absfun == NULL)
+			return INFINITY;
+		else
+			return absfun->ep.prob;
+	}
+	default:
+		return 0;
+	}
+}
diff --git a/src/runtime/c/pgf/expr.h b/src/runtime/c/pgf/expr.h
index a30e44318..7f8746b28 100644
--- a/src/runtime/c/pgf/expr.h
+++ b/src/runtime/c/pgf/expr.h
@@ -226,4 +226,7 @@ PGF_API_DECL void
 pgf_print_expr_tuple(size_t n_exprs, PgfExpr exprs[], PgfPrintContext* ctxt,
                      GuOut* out, GuExn* err);
 
+PGF_API prob_t
+pgf_compute_tree_probability(PgfPGF *gr, PgfExpr expr);
+
 #endif /* EXPR_H_ */