4 files changed, 76 insertions, 34 deletions
diff --git a/zap/source/amd64/mem/memcnt.S b/zap/source/amd64/mem/memcnt.S
new file mode 100644
index 0000000..5323301
--- /dev/null
+++ b/zap/source/amd64/mem/memcnt.S
@@ -0,0 +1,75 @@
+# Copyright 2022 Gabriel Jensen.
+# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
+# If a copy of the MPL was not distributed with this file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+.globl zap_memcnt
+
+zap_memcnt:
+# Address of the current element:
+#define addr     %rbx
+# Address of the element after the last input element:
+#define afterbuf %r12
+# Size of each input element:
+#define sz       %r13
+# Address of the function:
+#define fn       %r14
+# Count:
+#define cnt      %r15
+
+	# Push the callee-saved registers:
+	pushq addr
+	pushq afterbuf
+	pushq sz
+	pushq fn
+	pushq cnt
+
+	# Move registers into place:
+	movq %rdi,addr
+	movq %rsi,sz
+	movq %rdx,afterbuf
+	movq %rcx,fn
+
+	# Get the one-past-the-end address:
+	imulq sz,afterbuf  # afterbuf *= sz
+	addq addr,afterbuf # afterbuf += addr
+
+	movq $0x0,cnt # cnt = 0x0
+
+	# Iterate through the array:
+.loop:
+
+	# Check if we have reached the one-past-the-end address:
+	cmpq addr,afterbuf # if (addr == afterbuf)
+	je .done           # goto done
+
+	# Call the provided function:
+	movq addr,%rdi
+	call *fn       # fn(addr)
+
+	# Check the return value of the function:
+	testb %al,%al # if (fn(addr) == zap_false)
+	jz .skip      # goto skip
+
+	incq cnt # ++cnt
+
+	# Skip the incrementation:
+.skip:
+
+	# Continue to the next element:
+	addq sz,addr # addr += sz
+	jmp .loop    # goto loop
+
+	# Finish:
+.done:
+
+	# Move count into return register:
+	movq cnt,%rax
+	
+	# Restore the callee-saved registers:
+	popq cnt
+	popq fn
+	popq sz
+	popq afterbuf
+	popq addr
+
+	ret # return cnt
diff --git a/zap/source/amd64/mem/memcnt.c b/zap/source/amd64/mem/memcnt.c
deleted file mode 100644
index c22abca..0000000
--- a/zap/source/amd64/mem/memcnt.c
+++ /dev/null
@@ -1,22 +0,0 @@
-/*
-	Copyright 2022 Gabriel Jensen.
-	This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
-	If a copy of the MPL was not distributed with this file, You can obtain one at https://mozilla.org/MPL/2.0/.
-*/
-
-#include <zap/priv.h>
-
-#include <zap/mem.h>
-
-zap_sz zap_memcnt(void const * const _ptr,zap_sz const _sz,zap_sz const _num,zap_bool (* const _fn)(void const *)) {
-	unsigned char const *       addr     = _ptr;
-	zap_sz                const numbyte  = _num * _sz;
-	unsigned char const * const afterbuf = addr + numbyte;
-	zap_sz                      num      = 0x0u;
-	for (;addr != afterbuf;addr += _sz) {
-		if (_fn(addr)) {
-			++num;
-		}
-	}
-	return num;
-}
diff --git a/zap/source/amd64/mem/memfor.S b/zap/source/amd64/mem/memfor.S
index be705b9..e9cfe6b 100644
--- a/zap/source/amd64/mem/memfor.S
+++ b/zap/source/amd64/mem/memfor.S
@@ -25,10 +25,10 @@ zap_memfor:
 	# Move registers into place:
 	movq %rdi,addr
 	movq %rsi,sz
+	movq %rdx,afterbuf
 	movq %rcx,fn
 
 	# Get the one-past-the-end address:
-	movq %rdx,afterbuf
 	imulq sz,afterbuf  # afterbuf *= sz   // Calculate the array size in bytes (sz * num). We're using signed multiply because the equivalent using the unsigned instruction would use more instructions.
 	addq addr,afterbuf # afterbuf += addr
 
diff --git a/zap/source/amd64/mem/memgen.c b/zap/source/amd64/mem/memgen.c
index a39e326..5410bb3 100644
--- a/zap/source/amd64/mem/memgen.c
+++ b/zap/source/amd64/mem/memgen.c
@@ -6,17 +6,6 @@
 
 #include <zap/priv.h>
 
-/*.globl zap_memgen
-
-zap_memgen:*/
-
-	/*
-		void *  ptr
-		zap_sz  sz
-		zap_sz  num
-		void (* fn)(zap_sz,void *)
-	*/
-
 void zap_memgen(void * const _ptr,zap_sz const _sz,zap_sz const _num,void (* const _fn)(zap_sz,void *)) {
 	unsigned char *       ptr      = _ptr;
 	unsigned char * const afterbuf = ptr + _sz * _num;