added exprSize in the Haskell binding

4 files changed, 47 insertions, 2 deletions

diff --git a/src/runtime/c/pgf/expr.c b/src/runtime/c/pgf/expr.c
index f9fcd1442..751a7d25a 100644
--- a/src/runtime/c/pgf/expr.c
+++ b/src/runtime/c/pgf/expr.c
@@ -1177,6 +1177,39 @@ pgf_expr_hash(GuHash h, PgfExpr e)
 	return h;
 }
 
+PGF_API size_t
+pgf_expr_size(PgfExpr expr)
+{
+	GuVariantInfo ei = gu_variant_open(expr);
+	switch (ei.tag) {
+	case PGF_EXPR_ABS: {
+		PgfExprAbs* abs = ei.data;
+		return pgf_expr_size(abs->body);
+	}
+	case PGF_EXPR_APP: {
+		PgfExprApp* app = ei.data;
+		return pgf_expr_size(app->fun) + pgf_expr_size(app->arg);
+	}
+	case PGF_EXPR_LIT:
+	case PGF_EXPR_META:
+	case PGF_EXPR_FUN:
+	case PGF_EXPR_VAR: {
+		return 1;
+	}
+	case PGF_EXPR_TYPED: {
+		PgfExprTyped* typed = ei.data;
+		return pgf_expr_size(typed->expr);
+	}
+	case PGF_EXPR_IMPL_ARG: {
+		PgfExprImplArg* impl = ei.data;
+		return pgf_expr_size(impl->expr);
+	}
+	default:
+		gu_impossible();
+		return 0;
+	}
+}
+
 PGF_API void
 pgf_print_cid(PgfCId id,
 			  GuOut* out, GuExn* err)
diff --git a/src/runtime/c/pgf/expr.h b/src/runtime/c/pgf/expr.h
index 6492f8d18..962b3173f 100644
--- a/src/runtime/c/pgf/expr.h
+++ b/src/runtime/c/pgf/expr.h
@@ -197,6 +197,9 @@ pgf_literal_hash(GuHash h, PgfLiteral lit);
 PGF_API_DECL GuHash
 pgf_expr_hash(GuHash h, PgfExpr e);
 
+PGF_API size_t
+pgf_expr_size(PgfExpr expr);
+
 typedef struct PgfPrintContext PgfPrintContext;
 
 struct PgfPrintContext {
diff --git a/src/runtime/haskell-bind/PGF2.hsc b/src/runtime/haskell-bind/PGF2.hsc
index f2f1b4113..4990a1926 100644
--- a/src/runtime/haskell-bind/PGF2.hsc
+++ b/src/runtime/haskell-bind/PGF2.hsc
@@ -39,7 +39,7 @@ module PGF2 (-- * PGF
              mkFloat,unFloat,
              mkMeta,unMeta,
              mkCId,
-             exprHash,
+             exprHash, exprSize,
              treeProbability,
 
              -- ** Types
@@ -331,7 +331,13 @@ exprHash h (Expr c_expr touch1) =
     touch1
     return (fromIntegral h)
 
-    
+exprSize :: Expr -> Int32
+exprSize (Expr c_expr touch1) =
+  unsafePerformIO $ do
+    size <- pgf_expr_size c_expr
+    touch1
+    return (fromIntegral size)
+
 
 -----------------------------------------------------------------------------
 -- Graphviz
diff --git a/src/runtime/haskell-bind/PGF2/FFI.hs b/src/runtime/haskell-bind/PGF2/FFI.hs
index ce4dd6983..85c3a9793 100644
--- a/src/runtime/haskell-bind/PGF2/FFI.hs
+++ b/src/runtime/haskell-bind/PGF2/FFI.hs
@@ -331,6 +331,9 @@ foreign import ccall "pgf/expr.h pgf_expr_eq"
 foreign import ccall "pgf/expr.h pgf_expr_hash"
   pgf_expr_hash :: Word -> PgfExpr -> IO Word
 
+foreign import ccall "pgf/expr.h pgf_expr_size"
+  pgf_expr_size :: PgfExpr -> IO CInt
+
 foreign import ccall "pgf/expr.h pgf_compute_tree_probability"
   pgf_compute_tree_probability :: Ptr PgfPGF -> PgfExpr -> IO CFloat