rewrite malloc example

1 files changed, 6 insertions, 3 deletions

diff --git a/src/main.c b/src/main.c
index 9f55de1..0eb841c 100644
--- a/src/main.c
+++ b/src/main.c
@@ -92,7 +92,7 @@ is_punct(const char *p)
 		return 2;
 	if (strchr("-+", p[0]) != NULL && p[0] == p[1])
 		return 2;
-	if (p[0] == '|' && p[1] == '>')
+	if (strchr("|=", p[0]) != NULL && p[1] == '>')
 		return 2;
 
 	return (strchr("+-/*()<>,;{}[]=&^|", p[0]) != NULL);
@@ -679,7 +679,7 @@ stmt(Token **rest, Token *tok)
 		return loop_stmt(rest, tok);
 	if (equal(tok, "break"))
 		return break_stmt(rest, tok);
-	if (equal(tok, "return"))
+	if (equal(tok, "=>") || equal(tok, "return"))
 		return return_stmt(rest, tok);
 	if (equal(tok, "[")) {
 		int opn = 1;
@@ -742,7 +742,10 @@ expr_stmt(Token **rest, Token *tok)
 static Node *
 return_stmt(Token **rest, Token *tok)
 {
-	tok = skip(tok, "return");
+	if (tok->type == TOK_PUNCT)
+		tok = skip(tok, "=>");
+	else
+		tok = skip(tok, "return");
 	Node *node;
 	if (equal(tok, ";"))
 		node = new_unary(NOD_RETURN_STMT, new_num(0));