odinfmt fixes

2 files changed, 42 insertions, 16 deletions

diff --git a/src/odin/printer/printer.odin b/src/odin/printer/printer.odin
index 1241171..444148f 100644
--- a/src/odin/printer/printer.odin
+++ b/src/odin/printer/printer.odin
@@ -22,11 +22,16 @@ Printer :: struct {
 	indentation:          string,
 	newline:              string,
 	indentation_count:    int,
-	disabled_lines:       map[int]string,
-	last_disabled_line:   int,
+	disabled_lines:       map[int]Disabled_Info,
+	disabled_until_line:  int,
 	src:                  string,
 }
 
+Disabled_Info :: struct {
+	text: string,
+	end_line: int,
+}
+
 Config :: struct {
 	max_characters:       int,
 	spaces:               int,  //Spaces per indentation
@@ -110,8 +115,14 @@ build_disabled_lines_info :: proc(p: ^Printer) {
 			} else if strings.contains(comment.text[:], "//odinfmt: enable") && found_disable {
 				begin := disable_position.offset - (comment.pos.column - 1)
 				end := comment.pos.offset+len(comment.text)
+
+				disabled_info := Disabled_Info {
+					end_line = comment.pos.line,
+					text = p.src[begin:end],
+				}
+
 				for line := disable_position.line; line <= comment.pos.line; line += 1 {
-					p.disabled_lines[line] = p.src[begin:end]
+					p.disabled_lines[line] = disabled_info
 				}
 				
 				found_disable = false
diff --git a/src/odin/printer/visit.odin b/src/odin/printer/visit.odin
index b622ec7..b4c9f94 100644
--- a/src/odin/printer/visit.odin
+++ b/src/odin/printer/visit.odin
@@ -159,13 +159,18 @@ visit_comments :: proc(p: ^Printer, pos: tokenizer.Pos, end_newline := true) ->
 }
 
 visit_disabled :: proc(p: ^Printer, node: ^ast.Node) -> ^Document {
-	disabled_text := p.disabled_lines[node.pos.line]
+	
+	if node.pos.line not_in p.disabled_lines {
+		return empty()
+	}
+
+	disabled_info := p.disabled_lines[node.pos.line]
 
-	if p.last_disabled_line + 1 == node.pos.line {
-		return empty();
+	if disabled_info.text == "" {
+		return empty()
 	}
 
-	if disabled_text == "" {
+	if p.disabled_until_line > node.pos.line {
 		return empty()
 	}
 
@@ -175,11 +180,10 @@ visit_disabled :: proc(p: ^Printer, node: ^ast.Node) -> ^Document {
 		next_comment_group(p)
 	}
 
-	p.last_disabled_line = node.pos.line
-	
 	p.source_position = node.end
+	p.disabled_until_line = disabled_info.end_line
 
-	return cons(nest(-p.indentation_count, move), text(disabled_text))
+	return cons(nest(-p.indentation_count, move), text(disabled_info.text))
 }
 
 @(private)
@@ -856,6 +860,8 @@ visit_expr :: proc(p: ^Printer, expr: ^ast.Expr, called_from: Expr_Called_Type =
 		set_source_position(p, expr.end);
 	}
 
+	fmt.println(expr.derived)	
+
 	switch v in expr.derived {
 	case Inline_Asm_Expr:
 		document := cons(text_token(p, v.tok), text("("))
@@ -874,12 +880,21 @@ visit_expr :: proc(p: ^Printer, expr: ^ast.Expr, called_from: Expr_Called_Type =
 	case Auto_Cast:
 		return cons_with_nopl(text_token(p, v.op), visit_expr(p, v.expr))
 	case Ternary_If_Expr:
-		document := visit_expr(p, v.cond)
-		document = cons_with_nopl(document, text_token(p, v.op1))
-		document = cons_with_nopl(document, visit_expr(p, v.x))
-		document = cons_with_nopl(document, text_token(p, v.op2))
-		document = cons_with_nopl(document, visit_expr(p, v.y))
-		return document
+		if v.op1.text == "if" {
+			document := visit_expr(p, v.x)
+			document = cons_with_nopl(document, text_token(p, v.op1))
+			document = cons_with_nopl(document, visit_expr(p, v.cond))		
+			document = cons_with_nopl(document, text_token(p, v.op2))
+			document = cons_with_nopl(document, visit_expr(p, v.y))
+			return document
+		} else {
+			document := visit_expr(p, v.cond)
+			document = cons_with_nopl(document, text_token(p, v.op1))
+			document = cons_with_nopl(document, visit_expr(p, v.x))
+			document = cons_with_nopl(document, text_token(p, v.op2))
+			document = cons_with_nopl(document, visit_expr(p, v.y))
+			return document
+		}		
 	case Ternary_When_Expr:
 		document := visit_expr(p, v.cond)
 		document = cons_with_nopl(document, text_token(p, v.op1))