diff --git a/src/main/antlr4/PiccodeScript.g4 b/src/main/antlr4/PiccodeScript.g4
index 845a287..b14fef0 100644
--- a/src/main/antlr4/PiccodeScript.g4
+++ b/src/main/antlr4/PiccodeScript.g4
@@ -46,24 +46,15 @@ func_args: '(' arg_list? ')' ;
 
 arg_list: arg (',' arg)* ;
 
-arg: ID (ASSIGN literal_expr)? ; // only allow literals in default values
-
-
-literal_expr :
-	NUMBER                           
-	| STRING                           
-	|array
-	| tuple
-	| object
-	;
+arg: (USE?) ID (ASSIGN expr)? ; // only allow literals in default values
 
 expr_stmt: expr;
 	// parser rules
 expr
-	: expr CC expr
-	| expr LPAREN call_expr_list? RPAREN
+	: expr LPAREN call_expr_list? RPAREN
 	| var_decl
 	| closure_decl
+	| expr CC expr
 	| expr DOT expr
 	| expr MUL expr         
 	| expr DIV expr         
@@ -86,9 +77,9 @@ expr
 	| expr BAND expr         
 	| expr COLON expr
 	| LPAREN expr? RPAREN
-	| unary
 	| if_expr
 	| when_expr
+	| unary
 	| do_expr
 	| array
 	| tuple
@@ -102,12 +93,12 @@ closure_decl: BOR arg_list? BOR ARROW expr;
 unary: 
 	EXCLAIM expr
 	| SUB expr
-	| RETURN expr
+	| RETURN_TOK expr
 	| TILDE expr
 	| BAND expr;
 
 if_expr:
-	IF expr LBRACE expr RBRACE (ELSE LBRACE expr RBRACE)?;
+	IF expr expr (ELSE expr)?;
 
 when_expr: 
  WHEN expr LBRACE when_cases else_case? RBRACE;
@@ -182,6 +173,8 @@ IF: 'if';
 ELSE: 'else';
 MODULE: 'module';
 DO: 'do';
+USE: 'use';
+RETURN_TOK: 'return';
 
 NUMBER
     :   HEX_LITERAL
diff --git a/src/main/java/org/piccode/ast/PiccodeVisitor.java b/src/main/java/org/piccode/ast/PiccodeVisitor.java
index 4018238..fdcd46f 100644
--- a/src/main/java/org/piccode/ast/PiccodeVisitor.java
+++ b/src/main/java/org/piccode/ast/PiccodeVisitor.java
@@ -523,6 +523,15 @@ private Ast visitUnaryExpr(UnaryContext ctx) {
 			return result;
 		}
 
+
+		if (ctx.RETURN_TOK() != null) {
+			var tok = ctx.RETURN_TOK().getSymbol();
+			var result = finalizeAstNode(
+				new ReturnAst(visitExpr(ctx.expr())),
+				tok);
+			return result;
+		}
+
 		var start = ctx.getStart();
 		var err = new PiccodeException(fileName, start.getLine(), start.getCharPositionInLine(), "Invalid unary expression");
 		err.frame = null;
@@ -587,7 +596,20 @@ public Ast visitObject(ObjectContext ctx) {
 	private HashMap<String, Ast> visitKeyValuePairs(Key_val_pairsContext key_val_pairs) {
 		var obj = new HashMap<String, Ast>();
 		for (var kv : key_val_pairs.key_val_pair()) {
-			obj.put(kv.ID().getText(), visitExpr(kv.expr()));
+			var id = kv.ID();
+			if (id == null) {
+				var tok = key_val_pairs.getStart();
+				throw new PiccodeException(fileName, tok.getLine(), tok.getCharPositionInLine(), "Missing a key in object literal");
+			}
+
+			var expr = kv.expr();
+
+			if (expr == null) {
+				var tok = key_val_pairs.getStart();
+				throw new PiccodeException(fileName, tok.getLine(), tok.getCharPositionInLine(), "Missing expression in object literal");
+			}
+			
+			obj.put(id.getText(), visitExpr(expr));
 		}
 		return obj;
 	}