rewrite malloc example

2 files changed, 25 insertions, 17 deletions

diff --git a/samples/malloc.golem b/samples/malloc.golem
index 27502d4..a52d51e 100644
--- a/samples/malloc.golem
+++ b/samples/malloc.golem
@@ -1,18 +1,21 @@
-main() {
+main() => nopipe() |> pipe();
+
+nopipe() {
 	local s;
-	s = strdup("coucou le monde");
-	dbg s;
-	puts(s);
+	s = strdup("yo comment cava\n");
+	write(s);
 	stoupper(s);
-	puts(s);
+	write(s);
 }
 
-puts(s) {
+pipe() => strdup("coucou le monde\n") |> write() |> stoupper() |> write();
+
+write(s) {
+	local o;
+	o = s;
 	loop {
-		if ([s] == 0) {
-			wrt '\n';
-			return 0;
-		}
+		if ([s] == 0)
+			return o;
 		wrt [s];
 		s = s + 1;
 	}
@@ -22,9 +25,8 @@ strlen(s) {
 	local len;
 	len = 0;
 	loop {
-		if ([s + len] == 0) {
+		if ([s + len] == 0)
 			return len;
-		}
 		len = len + 1;
 	}
 }
@@ -45,8 +47,11 @@ strcpy(dst, src) {
 }
 
 stoupper(s) {
+	local o;
+	o = s;
 	loop {
-		if ([s] == 0) return 0;
+		if ([s] == 0)
+			return o;
 		[s] = [s] - ([s] >= 'a' & [s] <= 'z') * 32;
 		s = s + 1;
 	}
@@ -60,7 +65,7 @@ malloc(n) {
 	p = heap + heap_size;
 	heap_size = heap_size + n;
 	if (heap_size > 4096) {
-		puts("malloc error: heap is full");
+		write("malloc error: heap is full");
 		return 0;
 	}
 	return p;
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));