the reasoner's states are now closures just like the code for the def rules. This means more compact JIT compiled code and uniformity with the code for def rules

4 files changed, 85 insertions, 84 deletions

diff --git a/src/runtime/c/pgf/data.h b/src/runtime/c/pgf/data.h
index c2efed4e3..5b1842d74 100644
--- a/src/runtime/c/pgf/data.h
+++ b/src/runtime/c/pgf/data.h
@@ -82,7 +82,7 @@ typedef struct {
 	int arity;
 	PgfEquations* defns; // maybe null
 	PgfExprProb ep;
-	void* predicate;
+	PgfFunction predicate;
 	struct {
 		PgfFunction code;
 		union {
diff --git a/src/runtime/c/pgf/jit.c b/src/runtime/c/pgf/jit.c
index 53b1d8756..27a3752b4 100644
--- a/src/runtime/c/pgf/jit.c
+++ b/src/runtime/c/pgf/jit.c
@@ -154,9 +154,7 @@ pgf_jit_predicate(PgfReader* rdr, PgfAbstr* abstr,
 
 	pgf_jit_make_space(rdr, JIT_CODE_WINDOW);
 
-	abscat->predicate = (PgfPredicate) jit_get_ip().ptr;
-	
-	jit_prolog(2);
+	abscat->predicate = (PgfFunction) jit_get_ip().ptr;
 
 	PgfAbsFun* absfun = NULL;
 	PgfAbsFun* next_absfun = NULL;
@@ -170,17 +168,12 @@ pgf_jit_predicate(PgfReader* rdr, PgfAbstr* abstr,
 		gu_puts("\n", out, err);
 #endif
 
-		int rs_arg = jit_arg_p();
-		int parent_arg = jit_arg_p();
-		jit_getarg_p(JIT_V1, rs_arg);
-		jit_getarg_p(JIT_V2, parent_arg);
-
 		// compile TRY_FIRST
 		jit_prepare(3);
-		jit_movi_p(JIT_V0,next_absfun);
-		jit_pusharg_p(JIT_V0);
-		jit_pusharg_p(JIT_V2);
-		jit_pusharg_p(JIT_V1);
+		jit_movi_p(JIT_R0,next_absfun);
+		jit_pusharg_p(JIT_R0);
+		jit_pusharg_p(JIT_VCLOS);
+		jit_pusharg_p(JIT_VSTATE);
 		jit_finish(pgf_reasoner_try_first);
 	}
 
@@ -188,7 +181,7 @@ pgf_jit_predicate(PgfReader* rdr, PgfAbstr* abstr,
 	gu_puts("    RET\n", out, err);
 #endif
 	// compile RET
-	jit_ret();
+	jit_bare_ret();
 
 #ifdef PGF_JIT_DEBUG
 	if (n_funs > 0) {
@@ -201,13 +194,7 @@ pgf_jit_predicate(PgfReader* rdr, PgfAbstr* abstr,
 		pgf_jit_make_space(rdr, JIT_CODE_WINDOW);
 
 		absfun = next_absfun;
-		absfun->predicate = (PgfPredicate) jit_get_ip().ptr;
-
-		jit_prolog(2);
-		int rs_arg = jit_arg_p();
-		int st_arg = jit_arg_p();
-		jit_getarg_p(JIT_V1, rs_arg);
-		jit_getarg_p(JIT_V2, st_arg);
+		absfun->predicate = (PgfFunction) jit_get_ip().ptr;
 
 		size_t n_hypos = gu_seq_length(absfun->type->hypos);
 
@@ -223,13 +210,13 @@ pgf_jit_predicate(PgfReader* rdr, PgfAbstr* abstr,
 
 				// compile TRY_ELSE
 				jit_prepare(3);
-				jit_movi_p(JIT_V0, next_absfun);
-				jit_pusharg_p(JIT_V0);
-				jit_pusharg_p(JIT_V2);
-				jit_pusharg_p(JIT_V1);
+				jit_movi_p(JIT_R0, next_absfun);
+				jit_pusharg_p(JIT_R0);
+				jit_pusharg_p(JIT_VCLOS);
+				jit_pusharg_p(JIT_VSTATE);
 				jit_finish(pgf_reasoner_try_else);
 			}
-				
+
 			for (size_t i = 0; i < n_hypos; i++) {
 				PgfHypo* hypo = gu_seq_index(absfun->type->hypos, PgfHypo, i);
 
@@ -247,37 +234,23 @@ pgf_jit_predicate(PgfReader* rdr, PgfAbstr* abstr,
 #endif
 
 				// compile CALL
-				ref = jit_movi_p(JIT_V0, jit_forward());
-				jit_str_p(JIT_V2, JIT_V0);
-				jit_prepare(2);
-				jit_pusharg_p(JIT_V2);
-				jit_pusharg_p(JIT_V1);
+				ref = jit_movi_p(JIT_R0, jit_forward());
+				jit_str_p(JIT_VCLOS, JIT_R0);
 
 				PgfCallPatch patch;
 				patch.cid = hypo->type->cid;
-				patch.ref = jit_finish(jit_forward());
+				patch.ref = jit_jmpi(jit_forward());
 				gu_buf_push(rdr->jit_state->call_patches, PgfCallPatch, patch);
 
 #ifdef PGF_JIT_DEBUG
-				gu_puts("    RET\n", out, err);
 				if (i+1 < n_hypos) {
 					gu_printf(out, err, "L%d:\n", label++);
 				}
 #endif
 
-				// compile RET
-				jit_ret();
-
 				if (i+1 < n_hypos) {
 					pgf_jit_make_space(rdr, JIT_CODE_WINDOW);
-
 					jit_patch_movi(ref,jit_get_label());
-					
-					jit_prolog(2);
-					rs_arg = jit_arg_p();
-					st_arg = jit_arg_p();
-					jit_getarg_p(JIT_V1, rs_arg);
-					jit_getarg_p(JIT_V2, st_arg);
 				} else {
 					jit_patch_movi(ref,pgf_reasoner_complete);
 				}
@@ -294,10 +267,10 @@ pgf_jit_predicate(PgfReader* rdr, PgfAbstr* abstr,
 
 				// compile TRY_CONSTANT
 				jit_prepare(3);
-				jit_movi_p(JIT_V0, next_absfun);
-				jit_pusharg_p(JIT_V0);
-				jit_pusharg_p(JIT_V2);
-				jit_pusharg_p(JIT_V1);
+				jit_movi_p(JIT_R0, next_absfun);
+				jit_pusharg_p(JIT_R0);
+				jit_pusharg_p(JIT_VCLOS);
+				jit_pusharg_p(JIT_VSTATE);
 				jit_finish(pgf_reasoner_try_constant);
 			} else {
 #ifdef PGF_JIT_DEBUG
@@ -306,8 +279,8 @@ pgf_jit_predicate(PgfReader* rdr, PgfAbstr* abstr,
 
 				// compile COMPLETE
 				jit_prepare(2);
-				jit_pusharg_p(JIT_V2);
-				jit_pusharg_p(JIT_V1);
+				jit_pusharg_p(JIT_VCLOS);
+				jit_pusharg_p(JIT_VSTATE);
 				jit_finish(pgf_reasoner_complete);
 			}
 
@@ -382,9 +355,9 @@ pgf_jit_gates(PgfReader* rdr)
 
 	gates->enter = (void*) jit_get_ip().ptr;
 	jit_prolog(2);
-	int es_arg = jit_arg_p();
+	int rs_arg = jit_arg_p();
 	int closure_arg = jit_arg_p();
-	jit_getarg_p(JIT_VSTATE, es_arg);
+	jit_getarg_p(JIT_VSTATE, rs_arg);
 	jit_getarg_p(JIT_VCLOS,  closure_arg);
 	jit_stxi_p(offsetof(PgfReasoner, enter_stack_ptr), JIT_VSTATE, JIT_SP);
 	jit_ldr_p(JIT_R0, JIT_VCLOS);
@@ -614,6 +587,24 @@ pgf_jit_gates(PgfReader* rdr)
 	gates->evaluate_sum = jit_get_ip().ptr;
 	jit_jmpi(gates->mk_const);
 
+	pgf_jit_make_space(rdr, JIT_CODE_WINDOW);
+
+	gates->combine1 = (void*) jit_get_ip().ptr;
+	jit_prepare(2);
+	jit_pusharg_p(JIT_VCLOS);
+	jit_pusharg_p(JIT_VSTATE);
+	jit_finish(pgf_reasoner_combine1);
+	jit_bare_ret();
+
+	pgf_jit_make_space(rdr, JIT_CODE_WINDOW);
+
+	gates->combine2 = (void*) jit_get_ip().ptr;
+	jit_prepare(2);
+	jit_pusharg_p(JIT_VCLOS);
+	jit_pusharg_p(JIT_VSTATE);
+	jit_finish(pgf_reasoner_combine2);
+	jit_bare_ret();
+
 	gates->fin.fn = pgf_jit_finalize_cafs;
 	gates->cafs = NULL;
 	gu_pool_finally(rdr->opool, &gates->fin);
@@ -1291,7 +1282,7 @@ pgf_jit_done(PgfReader* rdr, PgfAbstr* abstr)
 		PgfAbsCat* arg =
 			gu_seq_binsearch(abstr->cats, pgf_abscat_order, PgfAbsCat, patch->cid);
 		if (arg != NULL) {
-			jit_patch_calli(patch->ref,(jit_insn*) arg->predicate);
+			jit_patch_at(patch->ref,(jit_insn*) arg->predicate);
 		} else {
 			PgfAbsFun* fun =
 				gu_seq_binsearch(abstr->funs, pgf_absfun_order, PgfAbsFun, patch->cid);
diff --git a/src/runtime/c/pgf/reasoner.c b/src/runtime/c/pgf/reasoner.c
index dfc8905ac..b53810c2a 100644
--- a/src/runtime/c/pgf/reasoner.c
+++ b/src/runtime/c/pgf/reasoner.c
@@ -18,8 +18,7 @@ typedef void (*PgfStatePrinter)(PgfReasonerState* st,
 #endif
 
 struct PgfReasonerState {
-	// the jitter expects that continuation is the first field
-	PgfPredicate continuation;
+	PgfClosure header;
 #ifdef PGF_REASONER_DEBUG
 	PgfStatePrinter print;
 #endif
@@ -129,7 +128,7 @@ pgf_combine1_to_expr(PgfCombine1State* st, GuPool* pool, GuPool* out_pool) {
 		gu_buf_get(st->exprs, PgfExprProb*, st->choice);
 
 	PgfExprState* nst = gu_new(PgfExprState, pool);
-	nst->base.continuation = st->parent->base.continuation;
+	nst->base.header.code = st->parent->base.header.code;
 	nst->base.prob = st->base.prob;
 	nst->answers = st->parent->answers;
 	nst->expr =
@@ -156,7 +155,7 @@ pgf_combine2_to_expr(PgfCombine2State* st, GuPool* pool, GuPool* out_pool)
 
 	PgfExprState* nst =
 		gu_new(PgfExprState, pool);
-	nst->base.continuation = parent->base.continuation;
+	nst->base.header.code = parent->base.header.code;
 	nst->base.prob         = st->base.prob;
 	nst->answers = parent->answers;
 	nst->expr    =
@@ -192,11 +191,12 @@ pgf_print_combine2_state(PgfCombine2State* st,
 }
 #endif
 
-static void
-pgf_combine1(PgfReasoner* rs, PgfCombine1State* st)
+void
+pgf_reasoner_combine1(PgfReasoner* rs, PgfClosure* closure)
 {
+	PgfCombine1State* st = (PgfCombine1State*) closure;
 	PgfExprState* nst = pgf_combine1_to_expr(st, rs->pool, rs->out_pool);
-	nst->base.continuation(rs, &nst->base);
+	rs->eval_gates->enter(rs, &nst->base.header);
 
 	st->choice++;
 
@@ -228,7 +228,7 @@ pgf_reasoner_try_first(PgfReasoner* rs, PgfExprState* parent, PgfAbsFun* absfun)
 
 	if (gu_buf_length(answers->parents) == 1) {
 		PgfExprState* st = gu_new(PgfExprState, rs->pool);
-		st->base.continuation = (PgfPredicate) absfun->predicate;
+		st->base.header.code = absfun->predicate;
 		st->base.prob = answers->outside_prob + absfun->ep.prob;
 		st->answers = answers;
 		st->expr = absfun->ep.expr;
@@ -245,7 +245,7 @@ pgf_reasoner_try_first(PgfReasoner* rs, PgfExprState* parent, PgfAbsFun* absfun)
 				gu_buf_get(answers->exprs, PgfExprProb*, 0);
 
 			PgfCombine1State* nst = gu_new(PgfCombine1State, rs->pool);
-			nst->base.continuation = (PgfPredicate) pgf_combine1;
+			nst->base.header.code  = rs->eval_gates->combine1;
 			nst->base.prob         = parent->base.prob + ep->prob;
 			nst->exprs             = answers->exprs;
 			nst->choice            = 0;
@@ -263,7 +263,7 @@ void
 pgf_reasoner_try_else(PgfReasoner* rs, PgfExprState* prev, PgfAbsFun* absfun)
 {
 	PgfExprState *st = gu_new(PgfExprState, rs->pool);
-	st->base.continuation = (PgfPredicate) absfun->predicate;
+	st->base.header.code = absfun->predicate;
 	st->base.prob = prev->answers->outside_prob + absfun->ep.prob;
 	st->answers   = prev->answers;
 	st->expr      = absfun->ep.expr;
@@ -275,12 +275,13 @@ pgf_reasoner_try_else(PgfReasoner* rs, PgfExprState* prev, PgfAbsFun* absfun)
 	gu_buf_heap_push(rs->pqueue, &pgf_expr_state_order, &st);
 }
 
-static void
-pgf_combine2(PgfReasoner* rs, PgfCombine2State* st)
+void
+pgf_reasoner_combine2(PgfReasoner* rs, PgfClosure* closure)
 {
+	PgfCombine2State* st = (PgfCombine2State*) closure;
 	PgfExprState* nst = pgf_combine2_to_expr(st, rs->pool, rs->out_pool);
 	if (nst != NULL) {
-		nst->base.continuation(rs, &nst->base);
+		rs->eval_gates->enter(rs, &nst->base.header);
 	}
 
 	st->choice++;
@@ -303,7 +304,7 @@ pgf_reasoner_complete(PgfReasoner* rs, PgfExprState* st)
 	gu_buf_push(st->answers->exprs, PgfExprProb*, ep);
 
 	PgfCombine2State* nst = gu_new(PgfCombine2State, rs->pool);
-	nst->base.continuation = (PgfPredicate) pgf_combine2;
+	nst->base.header.code  = rs->eval_gates->combine2;
 	nst->base.prob         = st->base.prob;
 	nst->parents           = st->answers->parents;
 	nst->choice            = 0;
@@ -312,7 +313,7 @@ pgf_reasoner_complete(PgfReasoner* rs, PgfExprState* st)
 #ifdef PGF_REASONER_DEBUG
 	nst->base.print = (PgfStatePrinter) pgf_print_combine2_state;
 #endif
-	nst->base.continuation(rs, &nst->base);
+	rs->eval_gates->enter(rs, &nst->base.header);
 }
 
 void
@@ -341,8 +342,8 @@ pgf_reasoner_next(PgfReasoner* rs)
 		}
 #endif
 
-		st->continuation(rs, st);
-		
+		rs->eval_gates->enter(rs, &st->header);
+
 		if (n_exprs < gu_buf_length(rs->exprs)) {
 			return gu_buf_get(rs->exprs, PgfExprProb*, n_exprs);
 		}
@@ -400,7 +401,9 @@ pgf_generate_all(PgfPGF* pgf, PgfCId cat, GuExn* err, GuPool* pool, GuPool* out_
 
 	PgfAbsCat* abscat = gu_seq_binsearch(rs->abstract->cats, pgf_abscat_order, PgfAbsCat, cat);
 	if (abscat != NULL) {
-		((PgfPredicate) abscat->predicate)(rs, NULL);
+		PgfClosure* closure = gu_new(PgfClosure, rs->pool);
+		closure->code = abscat->predicate;
+		rs->eval_gates->enter(rs, closure);
 	}
 
 	return &rs->en;
diff --git a/src/runtime/c/pgf/reasoner.h b/src/runtime/c/pgf/reasoner.h
index 966ba11a5..714b31b4a 100644
--- a/src/runtime/c/pgf/reasoner.h
+++ b/src/runtime/c/pgf/reasoner.h
@@ -7,20 +7,6 @@ typedef struct PgfReasoner PgfReasoner;
 typedef struct PgfReasonerState PgfReasonerState;
 typedef struct PgfExprState PgfExprState;
 
-typedef void (*PgfPredicate)(PgfReasoner* rs, PgfReasonerState* st);
-
-void
-pgf_reasoner_try_first(PgfReasoner* rs, PgfExprState* parent, PgfAbsFun* absfun);
-
-void
-pgf_reasoner_try_else(PgfReasoner* rs, PgfExprState* prev, PgfAbsFun* absfun);
-
-void
-pgf_reasoner_complete(PgfReasoner* rs, PgfExprState* st);
-
-void
-pgf_reasoner_try_constant(PgfReasoner* rs, PgfExprState* prev, PgfAbsFun* absfun);
-
 typedef struct {
 	PgfFunction code;
 } PgfClosure;
@@ -117,6 +103,9 @@ struct PgfEvalGates {
 
 	PgfFunction mk_const;
 
+	PgfFunction combine1;
+	PgfFunction combine2;
+
 	PgfClosure* (*enter)(PgfReasoner* rs, PgfClosure* closure);
 
 	GuFinalizer fin;
@@ -126,6 +115,24 @@ struct PgfEvalGates {
 PgfReasoner*
 pgf_new_reasoner(PgfPGF* pgf, GuExn* err, GuPool* pool, GuPool* out_pool);
 
+void
+pgf_reasoner_try_first(PgfReasoner* rs, PgfExprState* parent, PgfAbsFun* absfun);
+
+void
+pgf_reasoner_try_else(PgfReasoner* rs, PgfExprState* prev, PgfAbsFun* absfun);
+
+void
+pgf_reasoner_combine1(PgfReasoner* rs, PgfClosure* closure);
+
+void
+pgf_reasoner_combine2(PgfReasoner* rs, PgfClosure* closure);
+
+void
+pgf_reasoner_complete(PgfReasoner* rs, PgfExprState* st);
+
+void
+pgf_reasoner_try_constant(PgfReasoner* rs, PgfExprState* prev, PgfAbsFun* absfun);
+
 PgfClosure*
 pgf_evaluate_expr_thunk(PgfReasoner* rs, PgfExprThunk* thunk);