From fd363d267a741d9bb7bd82f07b13b387e3fdd168 Mon Sep 17 00:00:00 2001
From: =?utf8?q?J=C3=A9r=C3=A9mie=20Astor?= <fennecdjay@gmail.com>
Date: Wed, 21 Jul 2021 18:49:06 +0200
Subject: [PATCH] :art: Put memoize pragma inside function

---
 include/emit.h     |  1 -
 include/env/func.h |  1 +
 src/emit/emit.c    | 13 +++++++------
 src/emit/memoize.c |  2 +-
 src/parse/check.c  |  4 ++++
 tests/pp/pragma.gw | 11 +++++++----
 6 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/include/emit.h b/include/emit.h
index 64063f8a..34a287c3 100644
--- a/include/emit.h
+++ b/include/emit.h
@@ -33,7 +33,6 @@ struct EmitterInfo_ {
   char *         escape;
   VM_Code (*emit_code)(const Emitter);
   VM_Code  code;
-  uint16_t memoize;
   uint16_t unroll;
   uint16_t line;
   bool     debug;
diff --git a/include/env/func.h b/include/env/func.h
index 0550fbac..34f63624 100644
--- a/include/env/func.h
+++ b/include/env/func.h
@@ -21,6 +21,7 @@ struct Func_ {
   struct Map_      upvalues;
   float            inline_mult;
   uint16_t         weight;
+  uint16_t         memoize;
   uint16_t         ref;
   uint16_t         vt_index;
   ae_flag          flag;
diff --git a/src/emit/emit.c b/src/emit/emit.c
index 5f9091c6..c10167d6 100644
--- a/src/emit/emit.c
+++ b/src/emit/emit.c
@@ -2536,9 +2536,10 @@ ANN static m_bool emit_func_def_return(const Emitter emit) {
   CHECK_BB(emit_defers(emit));
   emit_return_pc(emit, val);
   vector_clear(&emit->code->stack_return);
-  if (emit->info->memoize && fflag(emit->env->func, fflag_pure)) {
+  const Func f = emit->env->func;
+  if (f->memoize && fflag(f, fflag_pure)) {
     const Instr instr = emit_add_instr(emit, MemoizeStore);
-    instr->m_val      = emit->env->func->def->stack_depth;
+    instr->m_val      = f->def->stack_depth;
   }
   return GW_OK;
 }
@@ -2568,7 +2569,7 @@ ANN static inline VM_Code _emit_func_def_code(const Emitter emit,
 
 ANN static VM_Code emit_func_def_code(const Emitter emit, const Func func) {
   const VM_Code code = _emit_func_def_code(emit, func);
-  if (emit->info->memoize && fflag(func, fflag_pure)) code->is_memoize = true;
+  if (func->memoize && fflag(func, fflag_pure)) code->is_memoize = true;
   code->ret_type = func->def->base->ret_type;
   return code;
 }
@@ -2636,7 +2637,8 @@ ANN static m_bool emit_memoize(const Emitter emit, const Func_Def fdef) {
 }
 
 ANN static m_bool emit_fdef(const Emitter emit, const Func_Def fdef) {
-  if (emit->info->memoize && fflag(fdef->base->func, fflag_pure))
+  const Func f = fdef->base->func;
+  if (f->memoize && fflag(f, fflag_pure))
     CHECK_BB(emit_memoize(emit, fdef));
   CHECK_BB(emit_func_def_body(emit, fdef));
   emit_func_def_return(emit);
@@ -2646,7 +2648,7 @@ ANN static m_bool emit_fdef(const Emitter emit, const Func_Def fdef) {
 ANN static void emit_fdef_finish(const Emitter emit, const Func_Def fdef) {
   const Func func = fdef->base->func;
   func->code      = emit_func_def_code(emit, func);
-  if (emit->info->memoize && fflag(func, fflag_pure))
+  if (func->memoize && fflag(func, fflag_pure))
     func->code->memoize = memoize_ini(emit, func);
 }
 
@@ -2775,7 +2777,6 @@ ANN static VM_Code emit_free_stack(const Emitter emit) {
 }
 
 ANN static inline void emit_clear(const Emitter emit) {
-  emit->info->memoize = 0;
   emit->info->unroll  = 0;
   emit->info->line    = 0;
   emit->this_offset   = 0;
diff --git a/src/emit/memoize.c b/src/emit/memoize.c
index 18bbbe4b..176b272f 100644
--- a/src/emit/memoize.c
+++ b/src/emit/memoize.c
@@ -20,7 +20,7 @@ Memoize memoize_ini(const Emitter emit, const Func f) {
   vector_init(&m->v);
   m->ret_sz = f->def->base->ret_type->size;
   m->arg_sz = f->def->stack_depth;
-  m->limit  = emit->info->memoize;
+  m->limit  = f->memoize;
   return m;
 }
 
diff --git a/src/parse/check.c b/src/parse/check.c
index 58952baa..9b95ebd7 100644
--- a/src/parse/check.c
+++ b/src/parse/check.c
@@ -1411,6 +1411,10 @@ ANN static m_bool check_stmt_match(const Env env, const Stmt_Match stmt) {
 
 ANN static m_bool check_stmt_pp(const Env env, const Stmt_PP stmt) {
   if (stmt->pp_type == ae_pp_include) env->name = stmt->data;
+  // check for memoization
+  if (env->func && stmt->pp_type == ae_pp_pragma &&
+      !strncmp(stmt->data, "memoize", strlen("memoize")))
+    env->func->memoize = strtol(stmt->data + 7, NULL, 10);
   return GW_OK;
 }
 
diff --git a/tests/pp/pragma.gw b/tests/pp/pragma.gw
index 76812b5b..59956d2f 100644
--- a/tests/pp/pragma.gw
+++ b/tests/pp/pragma.gw
@@ -1,9 +1,12 @@
-#pragma memoize 3
 fun int recursive_fib(int n) {
-    if (n < 2)
-        return n;
-    return recursive_fib(n - 2) + recursive_fib(n - 1);
+#pragma memoize 3
+
+  if (n < 2)
+      return n;
+
+  return recursive_fib(n - 2) + recursive_fib(n - 1);
 }
+
 <<< 6 => recursive_fib >>>;
 <<< 6 => recursive_fib >>>;
 <<< 12 => recursive_fib >>>;
-- 
2.43.0