5 files changed, 39 insertions, 11 deletions

diff --git a/src/exhaustive/exhaustive.c b/src/exhaustive/exhaustive.c
index 7ad4705..58d8305 100644
--- a/src/exhaustive/exhaustive.c
+++ b/src/exhaustive/exhaustive.c
@@ -124,12 +124,7 @@ static void exhaustive_ginit(gen_f *generators) {
 		} else if (cfg->random) {
 			generators[OFFSET_A] = &a_gen_random;
 			generators[OFFSET_B] = &b_gen_random;
-		} else {
-			generators[OFFSET_A] = &a_gen_input;
-			generators[OFFSET_B] = &b_gen_input;
-		}
-
-		if (cfg->koblitz) {
+		} else if (cfg->koblitz) {
 			switch (cfg->koblitz_value) {
 				case 0:
 					generators[OFFSET_A] = &a_gen_zero;
@@ -141,6 +136,9 @@ static void exhaustive_ginit(gen_f *generators) {
 					break;
 			}
 			generators[OFFSET_B] = &b_gen_one;
+		} else {
+			generators[OFFSET_A] = &a_gen_input;
+			generators[OFFSET_B] = &b_gen_input;
 		}
 
 		if (cfg->prime) {
@@ -149,6 +147,8 @@ static void exhaustive_ginit(gen_f *generators) {
 			generators[OFFSET_ORDER] = &order_gen_cofactor;
 		} else if (cfg->method == METHOD_ANOMALOUS) {
 			generators[OFFSET_ORDER] = &anomalous_gen_order;
+		} else if (cfg->koblitz){
+			generators[OFFSET_ORDER] = &order_gen_koblitz;
 		} else {
 			generators[OFFSET_ORDER] = &order_gen_any;
 		}
diff --git a/src/gen/order.c b/src/gen/order.c
index da1eb48..1f0d829 100644
--- a/src/gen/order.c
+++ b/src/gen/order.c
@@ -4,6 +4,7 @@
  */
 #include "order.h"
 #include "exhaustive/arg.h"
+#include "math/koblitz.h"
 #include "io/input.h"
 
 GENERATOR(order_gen_input) {
@@ -77,6 +78,23 @@ GENERATOR(order_gen_prime) {
 	}
 }
 
+GENERATOR(order_gen_koblitz) {
+	pari_sp ltop = avma;
+	GEN order = koblitz_get_order(cfg->bits, itou(curve->a));
+	if (order) {
+		curve->order = gerepilecopy(ltop, order);
+	} else {
+		GEN ord = ellff_get_card(curve->curve);
+		if (isclone(ord)) {
+			curve->order = gerepilecopy(ltop, ord);
+		} else {
+			avma = ltop;
+			curve->order = ord;
+		}
+	}
+	return 1;
+}
+
 CHECK(order_check_pohlig_hellman) {
 	HAS_ARG(args);
 	pari_sp ltop = avma;
diff --git a/src/gen/order.h b/src/gen/order.h
index 53e2250..45b5e3e 100644
--- a/src/gen/order.h
+++ b/src/gen/order.h
@@ -68,6 +68,16 @@ GENERATOR(order_gen_cofactor);
 GENERATOR(order_gen_prime);
 
 /**
+ * GENERATOR(gen_f)
+ *
+ * @param curve A curve_t being generated
+ * @param args unused
+ * @param state
+ * @return state diff
+ */
+GENERATOR(order_gen_koblitz);
+
+/**
  *
  * @param curve
  * @param args
diff --git a/src/math/koblitz.c b/src/math/koblitz.c
index 10467e3..b553cf6 100644
--- a/src/math/koblitz.c
+++ b/src/math/koblitz.c
@@ -56,14 +56,14 @@ static int compare_koblitz(const void *a, const void *b) {
 	}
 }
 
-const koblitz_t *koblitz_find(unsigned int m, unsigned int a) {
-	koblitz_t searched = {m, a, NULL};
+const koblitz_t *koblitz_find(unsigned long m, unsigned long a) {
+	koblitz_t searched = {(unsigned int) m, (unsigned int) a, NULL};
 	return (koblitz_t *)bsearch(&searched, koblitz_curves,
 	                            sizeof(koblitz_curves) / sizeof(koblitz_t),
 	                            sizeof(koblitz_t), &compare_koblitz);
 }
 
-GEN koblitz_get_order(unsigned int m, unsigned int a) {
+GEN koblitz_get_order(unsigned long m, unsigned long a) {
 	const koblitz_t *found = koblitz_find(m, a);
 	if (found) {
 		return strtoi(found->hex_order);
diff --git a/src/math/koblitz.h b/src/math/koblitz.h
index 8122c44..d2c62ed 100644
--- a/src/math/koblitz.h
+++ b/src/math/koblitz.h
@@ -27,7 +27,7 @@ bool koblitz_is_curve(const curve_t *curve);
  * @param a
  * @return
  */
-const koblitz_t *koblitz_find(unsigned int m, unsigned int a);
+const koblitz_t *koblitz_find(unsigned long m, unsigned long a);
 
 /**
  * @brief
@@ -35,6 +35,6 @@ const koblitz_t *koblitz_find(unsigned int m, unsigned int a);
  * @param a
  * @return
  */
-GEN koblitz_get_order(unsigned int m, unsigned int a);
+GEN koblitz_get_order(unsigned long m, unsigned long a);
 
 #endif  // ECGEN_MATH_KOBLITZ_H