/* SPDX-License-Identifier: GPL-2.0
 *
 *	Network memory
 *
 *	Author:	Mina Almasry <almasrymina@google.com>
 */

#ifndef _NET_NETMEM_H
#define _NET_NETMEM_H

#include <linux/dma-mapping.h>
#include <linux/mm.h>
#include <net/net_debug.h>


struct netmem_desc {
	unsigned long _flags;
	unsigned long pp_magic;
	struct page_pool *pp;
	unsigned long _pp_mapping_pad;
	unsigned long dma_addr;
	atomic_long_t pp_ref_count;
};

#define NETMEM_DESC_ASSERT_OFFSET(pg, desc)        \
	static_assert(offsetof(struct page, pg) == \
		      offsetof(struct netmem_desc, desc))
NETMEM_DESC_ASSERT_OFFSET(flags, _flags);
NETMEM_DESC_ASSERT_OFFSET(pp_magic, pp_magic);
NETMEM_DESC_ASSERT_OFFSET(pp, pp);
NETMEM_DESC_ASSERT_OFFSET(_pp_mapping_pad, _pp_mapping_pad);
NETMEM_DESC_ASSERT_OFFSET(dma_addr, dma_addr);
NETMEM_DESC_ASSERT_OFFSET(pp_ref_count, pp_ref_count);
#undef NETMEM_DESC_ASSERT_OFFSET


static_assert(sizeof(struct netmem_desc) <= offsetof(struct page, _refcount));



DECLARE_STATIC_KEY_FALSE(page_pool_mem_providers);


#define NET_IOV 0x01UL

enum net_iov_type {
	NET_IOV_DMABUF,
	NET_IOV_IOURING,
};


struct net_iov {
	union {
		struct netmem_desc desc;

		
		struct {
			unsigned long _flags;
			unsigned long pp_magic;
			struct page_pool *pp;
			unsigned long _pp_mapping_pad;
			unsigned long dma_addr;
			atomic_long_t pp_ref_count;
		};
	};
	struct net_iov_area *owner;
	enum net_iov_type type;
};

struct net_iov_area {
	
	struct net_iov *niovs;
	size_t num_niovs;

	
	unsigned long base_virtual;
};


#define NET_IOV_ASSERT_OFFSET(desc, iov)                    \
	static_assert(offsetof(struct netmem_desc, desc) == \
		      offsetof(struct net_iov, iov))
NET_IOV_ASSERT_OFFSET(_flags, _flags);
NET_IOV_ASSERT_OFFSET(pp_magic, pp_magic);
NET_IOV_ASSERT_OFFSET(pp, pp);
NET_IOV_ASSERT_OFFSET(_pp_mapping_pad, _pp_mapping_pad);
NET_IOV_ASSERT_OFFSET(dma_addr, dma_addr);
NET_IOV_ASSERT_OFFSET(pp_ref_count, pp_ref_count);
#undef NET_IOV_ASSERT_OFFSET

static inline struct net_iov_area *net_iov_owner(const struct net_iov *niov)
{
	return niov->owner;
}

static inline unsigned int net_iov_idx(const struct net_iov *niov)
{
	return niov - net_iov_owner(niov)->niovs;
}




typedef unsigned long __bitwise netmem_ref;

static inline bool netmem_is_net_iov(const netmem_ref netmem)
{
	return (__force unsigned long)netmem & NET_IOV;
}


static inline struct page *__netmem_to_page(netmem_ref netmem)
{
	return (__force struct page *)netmem;
}

static inline struct page *netmem_to_page(netmem_ref netmem)
{
	if (WARN_ON_ONCE(netmem_is_net_iov(netmem)))
		return NULL;

	return __netmem_to_page(netmem);
}

static inline struct net_iov *netmem_to_net_iov(netmem_ref netmem)
{
	if (netmem_is_net_iov(netmem))
		return (struct net_iov *)((__force unsigned long)netmem &
					  ~NET_IOV);

	DEBUG_NET_WARN_ON_ONCE(true);
	return NULL;
}

static inline netmem_ref net_iov_to_netmem(struct net_iov *niov)
{
	return (__force netmem_ref)((unsigned long)niov | NET_IOV);
}

#define page_to_netmem(p)	(_Generic((p),			\
	const struct page * :	(__force const netmem_ref)(p),	\
	struct page * :		(__force netmem_ref)(p)))


static inline netmem_ref virt_to_netmem(const void *data)
{
	return page_to_netmem(virt_to_page(data));
}

static inline int netmem_ref_count(netmem_ref netmem)
{
	
	if (netmem_is_net_iov(netmem))
		return 1;

	return page_ref_count(netmem_to_page(netmem));
}

static inline unsigned long netmem_pfn_trace(netmem_ref netmem)
{
	if (netmem_is_net_iov(netmem))
		return 0;

	return page_to_pfn(netmem_to_page(netmem));
}


#define __pp_page_to_nmdesc(p)	(_Generic((p),				\
	const struct page * :	(const struct netmem_desc *)(p),	\
	struct page * :		(struct netmem_desc *)(p)))


#define pp_page_to_nmdesc(p)						\
({									\
	DEBUG_NET_WARN_ON_ONCE(!page_pool_page_is_pp(p));		\
	__pp_page_to_nmdesc(p);						\
})


static inline struct netmem_desc *__netmem_to_nmdesc(netmem_ref netmem)
{
	return (__force struct netmem_desc *)netmem;
}


static inline struct netmem_desc *netmem_to_nmdesc(netmem_ref netmem)
{
	void *p = (void *)((__force unsigned long)netmem & ~NET_IOV);

	if (netmem_is_net_iov(netmem))
		return &((struct net_iov *)p)->desc;

	return __pp_page_to_nmdesc((struct page *)p);
}


static inline struct page_pool *__netmem_get_pp(netmem_ref netmem)
{
	return __netmem_to_nmdesc(netmem)->pp;
}

static inline struct page_pool *netmem_get_pp(netmem_ref netmem)
{
	return netmem_to_nmdesc(netmem)->pp;
}

static inline atomic_long_t *netmem_get_pp_ref_count_ref(netmem_ref netmem)
{
	return &netmem_to_nmdesc(netmem)->pp_ref_count;
}

static inline bool netmem_is_pref_nid(netmem_ref netmem, int pref_nid)
{
	
	if (netmem_is_net_iov(netmem))
		return true;

	return page_to_nid(netmem_to_page(netmem)) == pref_nid;
}

static inline netmem_ref netmem_compound_head(netmem_ref netmem)
{
	
	if (netmem_is_net_iov(netmem))
		return netmem;

	return page_to_netmem(compound_head(netmem_to_page(netmem)));
}


static inline void *__netmem_address(netmem_ref netmem)
{
	return page_address(__netmem_to_page(netmem));
}

static inline void *netmem_address(netmem_ref netmem)
{
	if (netmem_is_net_iov(netmem))
		return NULL;

	return __netmem_address(netmem);
}


static inline bool netmem_is_pfmemalloc(netmem_ref netmem)
{
	if (netmem_is_net_iov(netmem))
		return false;

	return page_is_pfmemalloc(netmem_to_page(netmem));
}

static inline unsigned long netmem_get_dma_addr(netmem_ref netmem)
{
	return netmem_to_nmdesc(netmem)->dma_addr;
}

#if defined(CONFIG_NET_DEVMEM)
static inline bool net_is_devmem_iov(const struct net_iov *niov)
{
	return niov->type == NET_IOV_DMABUF;
}
#else
static inline bool net_is_devmem_iov(const struct net_iov *niov)
{
	return false;
}
#endif

void __get_netmem(netmem_ref netmem);
void __put_netmem(netmem_ref netmem);

static __always_inline void get_netmem(netmem_ref netmem)
{
	if (netmem_is_net_iov(netmem))
		__get_netmem(netmem);
	else
		get_page(netmem_to_page(netmem));
}

static __always_inline void put_netmem(netmem_ref netmem)
{
	if (netmem_is_net_iov(netmem))
		__put_netmem(netmem);
	else
		put_page(netmem_to_page(netmem));
}

#define netmem_dma_unmap_addr_set(NETMEM, PTR, ADDR_NAME, VAL)   \
	do {                                                     \
		if (!netmem_is_net_iov(NETMEM))                  \
			dma_unmap_addr_set(PTR, ADDR_NAME, VAL); \
		else                                             \
			dma_unmap_addr_set(PTR, ADDR_NAME, 0);   \
	} while (0)

static inline void netmem_dma_unmap_page_attrs(struct device *dev,
					       dma_addr_t addr, size_t size,
					       enum dma_data_direction dir,
					       unsigned long attrs)
{
	if (!addr)
		return;

	dma_unmap_page_attrs(dev, addr, size, dir, attrs);
}

#endif 
