Handle func return

3 files changed, 20 insertions, 8 deletions

diff --git a/src/generator.c b/src/generator.c
index 7a960bf..3f9cc3b 100644
--- a/src/generator.c
+++ b/src/generator.c
@@ -91,16 +91,26 @@ void gen_const_declaration() {
   fun_display_table();
 }
 
-void gen_function_declaration(const char name[], int return_type,
-                              int nb_param) {
+Type gen_function_declaration(const char name[], int return_type, int nb_param) {
   fun_add(name, return_type, nb_param);
   fprintf(output, "\n%s:\npush rbp\nmov rbp,rsp\n", name);
+  return return_type;
 }
 
 void gen_function_end_declaration() {
   fprintf(output, "mov rsp, rbp\npop rbp\nret\n");
 }
 
+void gen_function_return(Type expect, Type actual) {
+  if (actual != expect) {
+    fprintf(stderr, "Return type mismatch at line %d.", lineno);
+    exit(1);
+  }
+
+  if (actual != VOID) fprintf(output, "pop rax\n");
+  gen_function_end_declaration();
+}
+
 int gen_function_call(const char name[], int nb_param) {
   Type return_type = fun_lookup(name, nb_param);
   fprintf(output, "call %s\n", name);
diff --git a/src/generator.h b/src/generator.h
index 939ab13..205fc01 100644
--- a/src/generator.h
+++ b/src/generator.h
@@ -19,8 +19,9 @@ FILE *output;
 void gen_prologue();
 void gen_prologue_continue(int *bss_done);
 void gen_const_declaration();
-void gen_function_declaration(const char name[], int return_type, int nb_param);
+Type gen_function_declaration(const char name[], int return_type, int nb_param);
 void gen_function_end_declaration();
+void gen_function_return(Type expect, Type actual);
 int gen_function_call(const char name[], int nb_param);
 void gen_declaration(const char name[], int type, Scope scope);
 void gen_check(const char name[], Scope scope);
diff --git a/src/tpc.y b/src/tpc.y
index e8f2e41..3bd3df4 100644
--- a/src/tpc.y
+++ b/src/tpc.y
@@ -20,6 +20,7 @@ extern int lineno;
 int yylex();
 void yyerror(char *);
 static Scope scope = GLOBAL;
+static Type return_type = VOID;
 static int bss_done = 0;
 static int num_label = 0;
 static int num_if = 0;
@@ -88,11 +89,11 @@ DeclFoncts:
 ;
 DeclFonct:
   EnTeteFonct { scope = LOCAL; }
-  Corps       { gen_function_end_declaration(); scope = GLOBAL; }
+  Corps       { gen_function_end_declaration(); scope = GLOBAL; return_type = VOID; }
 ;
 EnTeteFonct:
-  TYPE IDENT PrologueCont '(' Parametres ')' {gen_function_declaration($<ident>2, $<type>1, $5);}
-| VOID IDENT PrologueCont '(' Parametres ')' {gen_function_declaration($<ident>2, 2, $5);}
+  TYPE IDENT PrologueCont '(' Parametres ')' { return_type = gen_function_declaration($<ident>2, $<type>1, $5); }
+| VOID IDENT PrologueCont '(' Parametres ')' { return_type = gen_function_declaration($<ident>2, VOID, $5); }
 ;
 
 PrologueCont: {gen_prologue_continue(&bss_done);};
@@ -115,8 +116,8 @@ SuiteInstr:
 Instr:
   Exp ';'
 | ';'
-| RETURN Exp ';'
-| RETURN ';'
+| RETURN Exp ';'                 { gen_function_return(return_type, $<type>2); scope = GLOBAL; return_type = VOID; }
+| RETURN ';'                     { gen_function_return(return_type, VOID); scope = GLOBAL; return_type = VOID; }
 | READE '(' IDENT ')' ';'        { gen_reade($<ident>3); }
 | READC '(' IDENT ')' ';'        { gen_readc($<ident>3); }
 | PRINT '(' Exp ')' ';'          { gen_print($<type>3);}