/*
 * lrsc_tiny.c
 *		Two-core reproducer for lost LR/SC updates on RISC-V.
 *
 * Core A releases an LWLock as LWLockRelease() does (state = &lock->state = lock + 4; waiting: a byte on another
 * page, at the page offset of state + n):
 *
 *		sb		pending_wakeup, (waiting)
 *		lw		old_state, 4(lock)
 *		lw		head, 8(lock)
 *		and		desired_state, old_state, mask
 *	1:	lr.w	cur, (state)
 *		bne		cur, old_state, 2f
 *		sc.w	tmp, desired_state, (state)
 *		bnez	tmp, 1b
 *	2:
 *
 * while core B adds 1 to state.  Depending on n, A's SC might succeed although B's add came in between; the exit
 * status is then 1.
 */
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <unistd.h>

#define FLAGS		(7u << 28)		/* HAS_WAITERS, RELEASE_OK, LOCKED */
/* set by A's release in round r, kept by the CAS */
#define TAG(r)		(((r) % 0xff + 1) << 8)
#define TAG_MASK	(0xff << 8)
#define PAGE_SIZE	4096
#define NBUFFERS	512		/* 32 KiB of 64-byte BufferDescs */

struct shared
{
	uint32_t	done __attribute__((aligned(128)));
	uint8_t		pgproc[PAGE_SIZE] __attribute__((aligned(PAGE_SIZE)));	/* waiter->lwWaiting */
	uint8_t		buffers[NBUFFERS][64];
};

/* round r's word: at offset 40 as content_lock.state in a BufferDesc */
static uint32_t *
word(struct shared *sh, int r)
{
	return (uint32_t *) &sh->buffers[r % NBUFFERS][40];
}

/* LWLockWakeup()'s last store and final CAS: mark the waiter, clear FLAGS; returns 1 if B came in first */
static int
wakeup_unlock(uint32_t *state, uint8_t *waiting)
{
	char	   *lock = (char *) state - 4;	/* the LWLock: tranche, state, waiters */
	long		old_state, head, desired_state, cur, tmp;

	__asm__ __volatile__(
		/* waiter->lwWaiting = LW_WS_PENDING_WAKEUP */
		"	sb		%[pending_wakeup], 0(%[waiting])\n"
		/* old_state = pg_atomic_read_u32(&lock->state) */
		"	lw		%[old_state], 4(%[lock])\n"
		/* proclist_is_empty(&lock->waiters) */
		"	lw		%[head], 8(%[lock])\n"
		/* desired_state = old_state & ~(HAS_WAITERS | RELEASE_OK | LOCKED) */
		"	and		%[desired_state], %[old_state], %[mask]\n"
		/* pg_atomic_compare_exchange_u32(&lock->state, &old_state, ...) */
		"1:	lr.w	%[cur], (%[state])\n"
		"	bne		%[cur], %[old_state], 2f\n"
		"	sc.w	%[tmp], %[desired_state], (%[state])\n"
		"	bnez	%[tmp], 1b\n"
		"2:\n"
		: [old_state] "=&r"(old_state), [head] "=&r"(head), [cur] "=&r"(cur), [tmp] "=&r"(tmp),
		  [desired_state] "=&r"(desired_state)
		: [lock] "r"(lock), [state] "r"(state), [waiting] "r"(waiting), [pending_wakeup] "r"(2L),
		  [mask] "r"((long) ~FLAGS)
		: "memory");
	if (cur == old_state)
		return 0;

	/* B came in: retry from what the LR saw, as the while (true) loop does */
	uint32_t	o = (uint32_t) cur;

	while (!__atomic_compare_exchange_n(state, &o, o & ~FLAGS, 0, __ATOMIC_SEQ_CST, __ATOMIC_SEQ_CST))
		;
	return 1;
}

int
main(int argc, char **argv)
{
	int			n = 0, c, r;
	int			retried = 0, lost = 0;
	struct shared *sh = mmap(NULL, sizeof(*sh), PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0);

	if (sh == MAP_FAILED)
		return 2;
	while ((c = getopt(argc, argv, "n:")) != -1)
		if (c == 'n')
			n = atoi(optarg);
		else
			return 2;

	if (fork() == 0)
	{
		/*
		 * B: once A has released state, add 1; as many rounds as A.  Only A's release sets this round's TAG (its CAS
		 * keeps it), and a reused buffer still has an older one, so B's add comes after that store.
		 */
		for (r = 1; r <= 1000000; r++)
		{
			uint32_t   *state = word(sh, r);

			while ((__atomic_load_n(state, __ATOMIC_RELAXED) & TAG_MASK) != TAG(r))
				;
			__atomic_fetch_add(state, 1, __ATOMIC_SEQ_CST);
			__atomic_store_n(&sh->done, r, __ATOMIC_RELEASE);
		}
		_exit(0);
	}

	/* A: release state, clear FLAGS, check B's +1 */
	for (r = 1; r <= 1000000; r++)
	{
		uint32_t   *state = word(sh, r), v;

		/* vary A's timing relative to B */
		for (volatile uint32_t j = 0; j < (r >> 12) % 257; j++)
			;

		__atomic_store_n(state, FLAGS | TAG(r), __ATOMIC_RELAXED);

		retried += wakeup_unlock(state, sh->pgproc + (((uintptr_t) state + n) & (PAGE_SIZE - 1)));

		while (__atomic_load_n(&sh->done, __ATOMIC_ACQUIRE) != r)
			;
		v = __atomic_load_n(state, __ATOMIC_SEQ_CST);
		if (v != TAG(r) + 1 && lost++ < 10)
			printf("LOST UPDATE in round %d: value=0x%08x expected=0x%08x\n", r, v, TAG(r) + 1);
	}

	printf("page_offset=%d rounds=%d retried=%d lost=%d: %s\n", n, r - 1, retried, lost, lost ? "FAILED" : "OK");
	return lost != 0;
}