From 82a53000446b4b25fd290b15795495ddef8157e0 Mon Sep 17 00:00:00 2001
From: Swastik Baranwal <swstkbaranwal@gmail.com>
Date: Tue, 12 Jan 2021 09:08:12 +0530
Subject: [PATCH] cgen: allow sort with`<` and `>` op overloading (#8042)

---
 CHANGELOG.md                                  |  1 +
 vlib/v/gen/array.v                            | 53 +++++++++++--------
 .../sorting_by_different_criteria_test.v      |  6 +++
 3 files changed, 37 insertions(+), 23 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3a8ff79c2d..ccbf788884 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@
 - Array decompose: `[1, 2, 3]...` is now `...[1, 2, 3]`
 - Treating `enum` as `int` and operations on `enum` except `==` and `!=` are removed for strict type checking.
 - Support `[manualfree] fn f1(){}` and `[manualfree] module m1`, for functions doing their own memory management.
+- Allow usage of `<` and `>` operators for struct in `.sort` method for arrays, i.e. `arr.sort(a < b)`.
 
 ## V 0.2.1
 *30 Dec 2020*
diff --git a/vlib/v/gen/array.v b/vlib/v/gen/array.v
index 610298f4d6..1512a50640 100644
--- a/vlib/v/gen/array.v
+++ b/vlib/v/gen/array.v
@@ -252,33 +252,40 @@ fn (mut g Gen) gen_array_sort(node ast.CallExpr) {
 			// Variables `a` and `b` are used in the `.sort(a < b)` syntax, so we can reuse them
 			// when generating the function as long as the args are named the same.
 			g.definitions.writeln('int $compare_fn ($styp* a, $styp* b) {')
-			field_type := g.typ(infix_expr.left_type)
-			left_expr_str := g.write_expr_to_string(infix_expr.left).replace_once('.',
-				'->')
-			right_expr_str := g.write_expr_to_string(infix_expr.right).replace_once('.',
-				'->')
-			g.definitions.writeln('$field_type a_ = $left_expr_str;')
-			g.definitions.writeln('$field_type b_ = $right_expr_str;')
-			mut op1, mut op2 := '', ''
-			if infix_expr.left_type == table.string_type {
-				if is_reverse {
-					op1 = 'string_gt(a_, b_)'
-					op2 = 'string_lt(a_, b_)'
-				} else {
-					op1 = 'string_lt(a_, b_)'
-					op2 = 'string_gt(a_, b_)'
-				}
+			sym := g.table.get_type_symbol(typ)
+			if sym.has_method('<') && infix_expr.left.str().len == 1 {
+				g.definitions.writeln('\tif (${styp}__lt(*a, *b)) { return -1; } else { return 1; }}')
+			} else if sym.has_method('>') && infix_expr.left.str().len == 1 {
+				g.definitions.writeln('\tif (${styp}__gt(*a, *b)) { return -1; } else { return 1; }}')
 			} else {
-				if is_reverse {
-					op1 = 'a_ > b_'
-					op2 = 'a_ < b_'
+				field_type := g.typ(infix_expr.left_type)
+				left_expr_str := g.write_expr_to_string(infix_expr.left).replace_once('.',
+					'->')
+				right_expr_str := g.write_expr_to_string(infix_expr.right).replace_once('.',
+					'->')
+				g.definitions.writeln('$field_type a_ = $left_expr_str;')
+				g.definitions.writeln('$field_type b_ = $right_expr_str;')
+				mut op1, mut op2 := '', ''
+				if infix_expr.left_type == table.string_type {
+					if is_reverse {
+						op1 = 'string_gt(a_, b_)'
+						op2 = 'string_lt(a_, b_)'
+					} else {
+						op1 = 'string_lt(a_, b_)'
+						op2 = 'string_gt(a_, b_)'
+					}
 				} else {
-					op1 = 'a_ < b_'
-					op2 = 'a_ > b_'
+					if is_reverse {
+						op1 = 'a_ > b_'
+						op2 = 'a_ < b_'
+					} else {
+						op1 = 'a_ < b_'
+						op2 = 'a_ > b_'
+					}
 				}
+				g.definitions.writeln('if ($op1) return -1;')
+				g.definitions.writeln('if ($op2) return 1; return 0; }\n')
 			}
-			g.definitions.writeln('if ($op1) return -1;')
-			g.definitions.writeln('if ($op2) return 1; return 0; }\n')
 		}
 	}
 	if is_reverse && !compare_fn.ends_with('_reverse') {
diff --git a/vlib/v/tests/sorting_by_different_criteria_test.v b/vlib/v/tests/sorting_by_different_criteria_test.v
index 22dc6b1617..b40297d3a3 100644
--- a/vlib/v/tests/sorting_by_different_criteria_test.v
+++ b/vlib/v/tests/sorting_by_different_criteria_test.v
@@ -7,6 +7,10 @@ struct Parent {
 	name  string
 }
 
+fn (p Parent) < (p1 Parent) bool {
+    return p.name < p1.name
+}
+
 fn test_sorting_by_different_criteria_in_same_function() {
 	mut arr := [
 		Parent{Child{0.2}, 'def'},
@@ -20,4 +24,6 @@ fn test_sorting_by_different_criteria_in_same_function() {
 	// println(arr)
 	arr.sort(a.child.f < b.child.f)
 	assert arr[0].name == 'xyz'
+	arr.sort(a < b)
+	assert arr[0].name == 'abc'
 }