//go:build amd64 && go1.22
// +build amd64,go1.22

// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#include "go_asm.h"
#include "asm_amd64.h"
#include "textflag.h"

TEXT ·Count(SB), NOSPLIT, $0-40
#ifndef hasPOPCNT
	CMPB golang·org∕x∕sys∕cpu·X86+const_offsetX86HasPOPCNT(SB), $1
	JEQ  2(PC)
	JMP  ·countGeneric(SB)

#endif
	MOVQ b_base+0(FP), SI
	MOVQ b_len+8(FP), BX
	MOVB c+24(FP), AL
	LEAQ ret+32(FP), R8

	LEAL -65(AX), CX // Check if the byte is a ASCII letter
	CMPB CL, $25
	JLS  count_case  // Byte sought is a ASCII letter
	ADDL $-97, AX
	CMPB AL, $25
	JHI  count

count_case:
	MOVB c+24(FP), AL
	JMP  countbodyCase<>(SB)

count:
	MOVB c+24(FP), AL
	JMP  countbody<>(SB)

TEXT ·CountString(SB), NOSPLIT, $0-32
#ifndef hasPOPCNT
	CMPB golang·org∕x∕sys∕cpu·X86+const_offsetX86HasPOPCNT(SB), $1
	JEQ  2(PC)
	JMP  ·countGenericString(SB)

#endif
	MOVQ s_base+0(FP), SI
	MOVQ s_len+8(FP), BX
	MOVB c+16(FP), AL
	LEAQ ret+24(FP), R8

	LEAL -65(AX), CX // Check if the byte is a ASCII letter
	CMPB CL, $25
	JLS  count_case  // Byte sought is a ASCII letter
	ADDL $-97, AX
	CMPB AL, $25
	JHI  count

count_case:
	MOVB c+16(FP), AL
	JMP  countbodyCase<>(SB)

count:
	MOVB c+16(FP), AL
	JMP  countbody<>(SB)

// input:
//   SI: data
//   BX: data len
//   AL: byte sought
//   R8: address to put result
// This function requires the POPCNT instruction.
TEXT countbodyCase<>(SB), NOSPLIT, $0
	// Shuffle X0 around so that each byte contains
	// the character we're looking for.
	ORL       $32, AX    // Convert byte to lowercase
	MOVD      AX, X0
	PUNPCKLBW X0, X0
	PUNPCKLBW X0, X0
	PSHUFL    $0, X0, X0

	// Add space (' ') mask to X2
	MOVQ      $32, CX
	MOVQ      CX, X2
	PUNPCKLBW X2, X2
	PUNPCKLBW X2, X2
	PSHUFL    $0, X2, X2

	CMPQ BX, $16
	JLT  small

	MOVQ $0, R12 // Accumulator

	MOVQ SI, DI

	CMPQ BX, $64
	JA   avx2

sse:
	LEAQ -16(SI)(BX*1), AX // AX = address of last 16 bytes
	JMP  sseloopentry

	PCALIGN $16

sseloop:
	// Move the next 16-byte chunk of the data into X1.
	MOVOU (DI), X1

	// Logical OR to convert data to lowercase
	POR X2, X1

	// Compare bytes in X0 to X1.
	PCMPEQB X0, X1

	// Take the top bit of each byte in X1 and put the result in DX.
	PMOVMSKB X1, DX

	// Count number of matching bytes
	POPCNTL DX, DX

	// Accumulate into R12
	ADDQ DX, R12

	// Advance to next block.
	ADDQ $16, DI

sseloopentry:
	CMPQ DI, AX
	JBE  sseloop

	// Get the number of bytes to consider in the last 16 bytes
	ANDQ $15, BX
	JZ   end

	// Create mask to ignore overlap between previous 16 byte block
	// and the next.
	MOVQ $16, CX
	SUBQ BX, CX
	MOVQ $0xFFFF, R10
	SARQ CL, R10
	SALQ CL, R10

	// Process the last 16-byte chunk. This chunk may overlap with the
	// chunks we've already searched so we need to mask part of it.
	MOVOU    (AX), X1
	POR      X2, X1   // Convert data to lowercase
	PCMPEQB  X0, X1
	PMOVMSKB X1, DX

	// Apply mask
	ANDQ    R10, DX
	POPCNTL DX, DX
	ADDQ    DX, R12

end:
	MOVQ R12, (R8)
	RET

// handle for lengths < 16
small:
	TESTQ BX, BX
	JEQ   endzero

	// Check if we'll load across a page boundary.
	LEAQ  16(SI), AX
	TESTW $0xff0, AX
	JEQ   endofpage

	// We must ignore high bytes as they aren't part of our slice.
	// Create mask.
	MOVB BX, CX
	MOVQ $1, R10
	SALQ CL, R10
	SUBQ $1, R10

	// Load data
	MOVOU (SI), X1

	// Convert data to lowercase
	POR X2, X1

	// Compare target byte with each byte in data.
	PCMPEQB X0, X1

	// Move result bits to integer register.
	PMOVMSKB X1, DX

	// Apply mask
	ANDQ    R10, DX
	POPCNTL DX, DX

	// Directly return DX, we don't need to accumulate
	// since we have <16 bytes.
	MOVQ DX, (R8)
	RET

endzero:
	MOVQ $0, (R8)
	RET

endofpage:
	// We must ignore low bytes as they aren't part of our slice.
	MOVQ $16, CX
	SUBQ BX, CX
	MOVQ $0xFFFF, R10
	SARQ CL, R10
	SALQ CL, R10

	// Load data into the high end of X1.
	MOVOU -16(SI)(BX*1), X1

	// Convert data to lowercase
	POR X2, X1

	// Compare target byte with each byte in data.
	PCMPEQB X0, X1

	// Move result bits to integer register.
	PMOVMSKB X1, DX

	// Apply mask
	ANDQ R10, DX

	// Directly return DX, we don't need to accumulate
	// since we have <16 bytes.
	POPCNTL DX, DX
	MOVQ    DX, (R8)
	RET

avx2:
#ifndef hasAVX2
	CMPB golang·org∕x∕sys∕cpu·X86+const_offsetX86HasAVX2(SB), $1
	JNE  sse

#endif
	// Create a mask in Y6 that converts text to upper case.
	VPBROADCASTB X2, Y6
	MOVD         AX, X0
	LEAQ         -64(SI)(BX*1), R11
	LEAQ         (SI)(BX*1), R13
	VPBROADCASTB X0, Y1
	PCALIGN      $32

avx2_loop:
	VMOVDQU   (DI), Y2
	VMOVDQU   32(DI), Y4
	VPOR      Y6, Y2, Y2 // Convert data to lowercase
	VPOR      Y6, Y4, Y4 // Convert data to lowercase
	VPCMPEQB  Y1, Y2, Y3
	VPCMPEQB  Y1, Y4, Y5
	VPMOVMSKB Y3, DX
	VPMOVMSKB Y5, CX
	POPCNTL   DX, DX
	POPCNTL   CX, CX
	ADDQ      DX, R12
	ADDQ      CX, R12
	ADDQ      $64, DI
	CMPQ      DI, R11
	JLE       avx2_loop

	// If last block is already processed,
	// skip to the end.
	//
	// This check is NOT an optimization; if the input length is a
	// multiple of 64, we must not go through the last leg of the
	// function because the bit shift count passed to SALQ below would
	// be 64, which is outside of the 0-63 range supported by those
	// instructions.
	//
	// Tests in the bytes and strings packages with input lengths that
	// are multiples of 64 will break if this condition were removed.
	CMPQ DI, R13
	JEQ  endavx

	// Load address of the last 64 bytes.
	// There is an overlap with the previous block.
	MOVQ      R11, DI
	VMOVDQU   (DI), Y2
	VMOVDQU   32(DI), Y4
	VPOR      Y6, Y2, Y2 // Convert data to lowercase
	VPOR      Y6, Y4, Y4 // Convert data to lowercase
	VPCMPEQB  Y1, Y2, Y3
	VPCMPEQB  Y1, Y4, Y5
	VPMOVMSKB Y3, DX
	VPMOVMSKB Y5, CX

	// Exit AVX mode.
	VZEROUPPER
	SALQ $32, CX
	ORQ  CX, DX

	// Create mask to ignore overlap between previous 64 byte block
	// and the next.
	ANDQ $63, BX
	MOVQ $64, CX
	SUBQ BX, CX
	MOVQ $0xFFFFFFFFFFFFFFFF, R10
	SALQ CL, R10

	// Apply mask
	ANDQ    R10, DX
	POPCNTQ DX, DX
	ADDQ    DX, R12
	MOVQ    R12, (R8)
	RET

endavx:
	// Exit AVX mode.
	VZEROUPPER
	MOVQ R12, (R8)
	RET

// input:
//   SI: data
//   BX: data len
//   AL: byte sought
//   R8: address to put result
// This function requires the POPCNT instruction.
TEXT countbody<>(SB), NOSPLIT, $0
	// Shuffle X0 around so that each byte contains
	// the character we're looking for.
	MOVD      AX, X0
	PUNPCKLBW X0, X0
	PUNPCKLBW X0, X0
	PSHUFL    $0, X0, X0

	CMPQ BX, $16
	JLT  small

	MOVQ $0, R12 // Accumulator

	MOVQ SI, DI

	CMPQ BX, $64
	JAE  avx2

sse:
	LEAQ -16(SI)(BX*1), AX // AX = address of last 16 bytes
	JMP  sseloopentry

	PCALIGN $16

sseloop:
	// Move the next 16-byte chunk of the data into X1.
	MOVOU (DI), X1

	// Compare bytes in X0 to X1.
	PCMPEQB X0, X1

	// Take the top bit of each byte in X1 and put the result in DX.
	PMOVMSKB X1, DX

	// Count number of matching bytes
	POPCNTL DX, DX

	// Accumulate into R12
	ADDQ DX, R12

	// Advance to next block.
	ADDQ $16, DI

sseloopentry:
	CMPQ DI, AX
	JBE  sseloop

	// Get the number of bytes to consider in the last 16 bytes
	ANDQ $15, BX
	JZ   end

	// Create mask to ignore overlap between previous 16 byte block
	// and the next.
	MOVQ $16, CX
	SUBQ BX, CX
	MOVQ $0xFFFF, R10
	SARQ CL, R10
	SALQ CL, R10

	// Process the last 16-byte chunk. This chunk may overlap with the
	// chunks we've already searched so we need to mask part of it.
	MOVOU    (AX), X1
	PCMPEQB  X0, X1
	PMOVMSKB X1, DX

	// Apply mask
	ANDQ    R10, DX
	POPCNTL DX, DX
	ADDQ    DX, R12

end:
	MOVQ R12, (R8)
	RET

// handle for lengths < 16
small:
	TESTQ BX, BX
	JEQ   endzero

	// Check if we'll load across a page boundary.
	LEAQ  16(SI), AX
	TESTW $0xff0, AX
	JEQ   endofpage

	// We must ignore high bytes as they aren't part of our slice.
	// Create mask.
	MOVB BX, CX
	MOVQ $1, R10
	SALQ CL, R10
	SUBQ $1, R10

	// Load data
	MOVOU (SI), X1

	// Compare target byte with each byte in data.
	PCMPEQB X0, X1

	// Move result bits to integer register.
	PMOVMSKB X1, DX

	// Apply mask
	ANDQ    R10, DX
	POPCNTL DX, DX

	// Directly return DX, we don't need to accumulate
	// since we have <16 bytes.
	MOVQ DX, (R8)
	RET

endzero:
	MOVQ $0, (R8)
	RET

endofpage:
	// We must ignore low bytes as they aren't part of our slice.
	MOVQ $16, CX
	SUBQ BX, CX
	MOVQ $0xFFFF, R10
	SARQ CL, R10
	SALQ CL, R10

	// Load data into the high end of X1.
	MOVOU -16(SI)(BX*1), X1

	// Compare target byte with each byte in data.
	PCMPEQB X0, X1

	// Move result bits to integer register.
	PMOVMSKB X1, DX

	// Apply mask
	ANDQ R10, DX

	// Directly return DX, we don't need to accumulate
	// since we have <16 bytes.
	POPCNTL DX, DX
	MOVQ    DX, (R8)
	RET

avx2:
#ifndef hasAVX2
	CMPB golang·org∕x∕sys∕cpu·X86+const_offsetX86HasAVX2(SB), $1
	JNE  sse

#endif
	MOVD         AX, X0
	LEAQ         -64(SI)(BX*1), R11
	LEAQ         (SI)(BX*1), R13
	VPBROADCASTB X0, Y1
	PCALIGN      $32

avx2_loop:
	VMOVDQU   (DI), Y2
	VMOVDQU   32(DI), Y4
	VPCMPEQB  Y1, Y2, Y3
	VPCMPEQB  Y1, Y4, Y5
	VPMOVMSKB Y3, DX
	VPMOVMSKB Y5, CX
	POPCNTL   DX, DX
	POPCNTL   CX, CX
	ADDQ      DX, R12
	ADDQ      CX, R12
	ADDQ      $64, DI
	CMPQ      DI, R11
	JLE       avx2_loop

	// If last block is already processed,
	// skip to the end.
	//
	// This check is NOT an optimization; if the input length is a
	// multiple of 64, we must not go through the last leg of the
	// function because the bit shift count passed to SALQ below would
	// be 64, which is outside of the 0-63 range supported by those
	// instructions.
	//
	// Tests in the bytes and strings packages with input lengths that
	// are multiples of 64 will break if this condition were removed.
	CMPQ DI, R13
	JEQ  endavx

	// Load address of the last 64 bytes.
	// There is an overlap with the previous block.
	MOVQ      R11, DI
	VMOVDQU   (DI), Y2
	VMOVDQU   32(DI), Y4
	VPCMPEQB  Y1, Y2, Y3
	VPCMPEQB  Y1, Y4, Y5
	VPMOVMSKB Y3, DX
	VPMOVMSKB Y5, CX

	// Exit AVX mode.
	VZEROUPPER
	SALQ $32, CX
	ORQ  CX, DX

	// Create mask to ignore overlap between previous 64 byte block
	// and the next.
	ANDQ $63, BX
	MOVQ $64, CX
	SUBQ BX, CX
	MOVQ $0xFFFFFFFFFFFFFFFF, R10
	SALQ CL, R10

	// Apply mask
	ANDQ    R10, DX
	POPCNTQ DX, DX
	ADDQ    DX, R12
	MOVQ    R12, (R8)
	RET

endavx:
	// Exit AVX mode.
	VZEROUPPER
	MOVQ R12, (R8)
	RET
