/*
 * Copyright (C) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */


#include <stdio.h>
#include <malloc.h>
#include <stdlib.h>
#include <pthread.h>
#include <stdint.h>
#include <memory.h>

#include "test-malloc-api-common.h"

#define BARRIER_HEIGHT 2
#define ALLOCATIONS_NUMBER 8
#define MIN(x, y) (((x) < (y)) ? (x) : (y))

typedef struct iterate_arg_s {
	uintptr_t allocs[ALLOCATIONS_NUMBER];
	size_t allocs_reported_number[ALLOCATIONS_NUMBER];
	size_t allocs_actual_sizes[ALLOCATIONS_NUMBER];
	size_t reported_sizes[ALLOCATIONS_NUMBER];
} iterate_arg_t;

typedef struct {
	uintptr_t *base;
	size_t size;
} allocations_info_t;

static const size_t allocs_sizes[ALLOCATIONS_NUMBER] = {
	8,
	2 * 1024,
	64 * 1024,
	512 * 1024,
	2 * 1024 * 1024,
	8 * 1024 * 1024,
	16 * 1024 * 1024,
	32 * 1024 * 1024
};

void iterate_callback(void *base, size_t size, void *data)
{
	iterate_arg_t *iterate_arg = (iterate_arg_t *) data;
	uintptr_t end;
	if (__builtin_add_overflow((uintptr_t) base, size, &end)) {
		return;
	}

	for (size_t i = 0; i < ALLOCATIONS_NUMBER; ++i) {
		if (iterate_arg->allocs[i] >= (uintptr_t) base && iterate_arg->allocs[i] < end) {
			iterate_arg->allocs_reported_number[i]++;
			uintptr_t max_size = end - iterate_arg->allocs[i];
			iterate_arg->reported_sizes[i] = MIN(size, max_size);
		}
	}
}

void fill_allocations_info(const iterate_arg_t *iterate_arg, allocations_info_t *allocations_info)
{
	size_t min_idx, max_idx;
	uintptr_t min_val = UINTPTR_MAX, max_val = 0;

	const uintptr_t *allocs = iterate_arg->allocs;

	for (size_t i = 0; i < ALLOCATIONS_NUMBER; ++i) {
		if (allocs[i] > max_val) {
			max_val = allocs[i];
			max_idx = i;
		}
		if (allocs[i] < min_val) {
			min_val = allocs[i];
			min_idx = i;
		}
	}

	allocations_info->base = (void *) allocs[min_idx];
	allocations_info->size = allocs[max_idx] - allocs[min_idx] + allocs_sizes[max_idx];
}

void make_allocations(iterate_arg_t *iterate_arg)
{
	uintptr_t *allocs = iterate_arg->allocs;
	size_t *allocs_actual_sizes = iterate_arg->allocs_actual_sizes;

	for (size_t i = 0; i < ALLOCATIONS_NUMBER; ++i) {
		allocs[i] = (uintptr_t) malloc(allocs_sizes[i]);
		allocs_actual_sizes[i] = malloc_usable_size((void *) allocs[i]);
	}
}

void free_allocations(iterate_arg_t *iterate_arg)
{
	uintptr_t *allocs = iterate_arg->allocs;

	for (size_t i = 0; i < ALLOCATIONS_NUMBER; ++i) {
		free((void *) allocs[i]);
	}
}

int iterate_wrapper(iterate_arg_t *iterate_arg)
{
	int ret = 0;
	allocations_info_t allocations_info;
	fill_allocations_info(iterate_arg, &allocations_info);
	malloc_iterate(allocations_info.base, allocations_info.size, iterate_callback, iterate_arg);

	for (size_t i = 0; i < ALLOCATIONS_NUMBER; ++i) {
		if (iterate_arg->allocs_reported_number[i] != 1) {
			ret = -1;
		}
	}
	return ret;
}

pthread_barrier_t routine_allocated;
pthread_barrier_t routine_iterated;

void *allocate_routine(void *vargp)
{
	iterate_arg_t *iterate_arg = (iterate_arg_t *) vargp;
	make_allocations(iterate_arg);
	pthread_barrier_wait(&routine_allocated);
	pthread_barrier_wait(&routine_iterated);
	return NULL;
}

void *abandoned_allocate_routine(void *vargp)
{
	iterate_arg_t *iterate_arg = (iterate_arg_t *) vargp;
	make_allocations(iterate_arg);
	return NULL;
}

int test_iterate_main_thread(void)
{
	int ret;
	iterate_arg_t iterate_arg = {{0}, {0}, {0}, {0}};
	make_allocations(&iterate_arg);
	ret = iterate_wrapper(&iterate_arg);
	free_allocations(&iterate_arg);
	return ret;
}

int test_iterate_another_thread(void)
{
	int ret;
	iterate_arg_t iterate_arg_routine = {{0}, {0}, {0}, {0}};
	pthread_barrier_init(&routine_allocated, NULL, BARRIER_HEIGHT);
	pthread_barrier_init(&routine_iterated, NULL, BARRIER_HEIGHT);
	pthread_t thread_id;
	pthread_create(&thread_id, NULL, allocate_routine, (void *) &iterate_arg_routine);
	pthread_barrier_wait(&routine_allocated);
	ret = iterate_wrapper(&iterate_arg_routine);
	free_allocations(&iterate_arg_routine);
	pthread_barrier_wait(&routine_iterated);
	return ret;
}

int test_iterate_over_abandoned_allocs(void)
{
	int ret;
	iterate_arg_t iterate_arg_routine = {{0}, {0}, {0}, {0}};
	pthread_t thread_id;
	pthread_create(&thread_id, NULL, abandoned_allocate_routine, (void *) &iterate_arg_routine);
	pthread_join(thread_id, NULL);
	ret = iterate_wrapper(&iterate_arg_routine);
	free_allocations(&iterate_arg_routine);
	return ret;
}

int main()
{
	int ret = 0;

	ret = check_and_report("Testing iterate main thread", test_iterate_main_thread);

	ret = -(ret || check_and_report("Testing iterate another thread", test_iterate_another_thread));

	ret = -(ret || check_and_report("Testing iterate over abandoned allocations", test_iterate_over_abandoned_allocs));

	return ret;
}