added function treeProbability in the Haskell binding

4 files changed, 41 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_ */
diff --git a/src/runtime/haskell-bind/PGF2.hsc b/src/runtime/haskell-bind/PGF2.hsc
index 4523279dd..54c413a34 100644
--- a/src/runtime/haskell-bind/PGF2.hsc
+++ b/src/runtime/haskell-bind/PGF2.hsc
@@ -39,6 +39,8 @@ module PGF2 (-- * PGF
              mkFloat,unFloat,
              mkMeta,unMeta,
              mkCId,
+             treeProbability,
+
              -- ** Types
              Type, Hypo, BindType(..), startCat,
              readType, showType,
@@ -314,6 +316,13 @@ compute (PGF p _) (Expr c_expr touch1) =
               gu_pool_free exprPl
               throwIO (PGFError msg)
 
+treeProbability :: PGF -> Expr -> Float
+treeProbability (PGF p _) (Expr c_expr touch1) =
+  unsafePerformIO $ do
+    res <- pgf_compute_tree_probability p c_expr
+    touch1
+    return (realToFrac res)
+
 -----------------------------------------------------------------------------
 -- Graphviz
 
diff --git a/src/runtime/haskell-bind/PGF2/FFI.hs b/src/runtime/haskell-bind/PGF2/FFI.hs
index 1a5e7f91b..65dd81085 100644
--- a/src/runtime/haskell-bind/PGF2/FFI.hs
+++ b/src/runtime/haskell-bind/PGF2/FFI.hs
@@ -325,6 +325,9 @@ foreign import ccall "pgf/pgf.h pgf_expr_unlit"
 foreign import ccall "pgf/expr.h pgf_expr_arity"
   pgf_expr_arity :: PgfExpr -> IO CInt
 
+foreign import ccall "pgf/expr.h pgf_compute_tree_probability"
+  pgf_compute_tree_probability :: Ptr PgfPGF -> PgfExpr -> IO CFloat
+
 foreign import ccall "pgf/expr.h pgf_check_expr"
   pgf_check_expr :: Ptr PgfPGF -> Ptr PgfExpr -> PgfType -> Ptr GuExn -> Ptr GuPool -> IO ()