sparse_type.h

/*
	Copyright (C) 2025-2026 Zhenjie Li (Li, Zhenjie), Xiang Li (Li, Xiang)

	This file is part of SparseRREF. The SparseRREF is free software:
	you can redistribute it and/or modify it under the terms of the MIT
	License.
*/


#ifndef SPARSE_TYPE_H
#define SPARSE_TYPE_H

#include <limits>
#include <numeric>

#include "sparse_rref.h"
#include "scalar.h"

namespace SparseRREF {
	enum SPARSE_TYPE {
		SPARSE_CSR, // Compressed sparse row
		SPARSE_COO, // Coordinate list
		SPARSE_LR,   // List of rows
		SPARSE_OTHER  // Other types
	};

	template <Flint::builtin_integral index_t = int>
	struct pivot_t {
		index_t r;
		index_t c;
	};

	// sparse vector
	template <typename T, Flint::builtin_integral index_t = int> struct sparse_vec {
		index_t* indices = nullptr;
		T* entries = nullptr;
		size_t _nnz = 0;
		size_t _alloc = 0;

		struct de_iterator_ref {
			index_t& ind;
			T& val;
		};

		struct de_const_iterator_ref {
			const index_t& ind;
			const T& val;
		};

		// the value type must stay dependent on the pointer type, so that a const sparse_vec
		// hands out references that cannot be used to modify it
		template <typename ind_p, typename val_p, typename ref> struct basic_iterator {
			ind_p ind_ptr;
			val_p val_ptr;

			basic_iterator& operator++() { ind_ptr++; val_ptr++; return *this; }
			basic_iterator operator++(int) { basic_iterator tmp = *this; ind_ptr++; val_ptr++; return tmp; }
			basic_iterator& operator--() { ind_ptr--; val_ptr--; return *this; }
			basic_iterator operator--(int) { basic_iterator tmp = *this; ind_ptr--; val_ptr--; return tmp; }
			basic_iterator& operator+=(size_t n) { ind_ptr += n; val_ptr += n; return *this; }
			basic_iterator& operator-=(size_t n) { ind_ptr -= n; val_ptr -= n; return *this; }
			basic_iterator operator+(size_t n) const { basic_iterator tmp = *this; tmp += n; return tmp; }
			basic_iterator operator-(size_t n) const { basic_iterator tmp = *this; tmp -= n; return tmp; }
			bool operator==(const basic_iterator& other) const { return ind_ptr == other.ind_ptr; }
			bool operator!=(const basic_iterator& other) const { return ind_ptr != other.ind_ptr; }

			ref operator*() const { return { *ind_ptr, *val_ptr }; }
		};

		using iterator = basic_iterator<index_t*, T*, de_iterator_ref>;
		using const_iterator = basic_iterator<const index_t*, const T*, de_const_iterator_ref>;

		// functions of iterator 
		iterator begin() { return { indices, entries }; }
		iterator end() { return { indices + _nnz, entries + _nnz }; }
		const_iterator begin() const { return { indices, entries }; }
		const_iterator end() const { return { indices + _nnz, entries + _nnz }; }
		const_iterator cbegin() const { return { indices, entries }; }
		const_iterator cend() const { return { indices + _nnz, entries + _nnz }; }

		std::span<index_t> index_span() { return std::span<index_t>(indices, _nnz); }
		std::span<const index_t> index_span() const { return std::span<const index_t>(indices, _nnz); }
		std::span<T> entry_span() { return std::span<T>(entries, _nnz); }
		std::span<const T> entry_span() const { return std::span<const T>(entries, _nnz); }

		// C++23 is needed for zip_view
		// auto index_view() const { return std::ranges::subrange(indices, indices + _nnz); }
		// auto entry_view() const { return std::ranges::subrange(entries, entries + _nnz); }
		// auto combine_view() const { return std::ranges::zip_view(index_view(), entry_view()); }

		sparse_vec() {
			indices = nullptr;
			entries = nullptr;
			_nnz = 0;
			_alloc = 0;
		}

		void clear() {
			if (_alloc == 0)
				return;
			s_free(indices);
			indices = nullptr;
			for (size_t i = 0; i < _alloc; i++)
				entries[i].~T();
			s_free(entries);
			entries = nullptr;
			_alloc = 0;
			_nnz = 0;
		}

		~sparse_vec() {
			clear();
		}

		void reserve(const size_t n, const bool is_copy = true) {
			if (n == _alloc)
				return;

			if (n == 0) {
				clear();
				return;
			}

			if (_alloc == 0) {
				indices = s_malloc<index_t>(n);
				entries = s_malloc<T>(n);
				for (size_t i = 0; i < n; i++)
					new (entries + i) T();
				_alloc = n;
				return;
			}

			// when expanding or using realloc, sometimes we need to copy the old data
			// if is_copy is false, we do not make sure that the old data is copied to 
			// the new memory, it is useful when we just want to enlarge the memory
			if (!is_copy && n > _alloc) {
				auto ii = s_expand(indices, n);
				auto ee = s_expand(entries, n);
				if (ii == nullptr) {
					s_free(indices);
					indices = s_malloc<index_t>(n);
				}
				else {
					indices = ii;
				}
				if (ee == nullptr) {
					for (size_t i = 0; i < _alloc; i++) {
						entries[i].~T();
					}
					s_free(entries);
					entries = s_malloc<T>(n);
					for (size_t i = 0; i < n; i++) {
						new (entries + i) T();
					}
				}
				else {
					entries = ee;
					for (size_t i = _alloc; i < n; i++) {
						new (entries + i) T();
					}
				}

				_alloc = n;
				_nnz = 0;
				return;
			}

			indices = s_realloc(indices, n);

			if (n < _alloc) {
				for (size_t i = n; i < _alloc; i++)
					entries[i].~T();
				entries = s_realloc<T>(entries, n);
			}
			else {
				entries = s_realloc<T>(entries, n);
				for (size_t i = _alloc; i < n; i++)
					new (entries + i) T();
			}

			_alloc = n;
		}

		inline void zero() { _nnz = 0; }
		inline void resize(const size_t n) { _nnz = n; }

		inline void copy(const sparse_vec& l) {
			if (this == &l)
				return;
			if (_alloc < l._nnz)
				reserve(l._nnz);
			for (size_t i = 0; i < l._nnz; i++) {
				indices[i] = l.indices[i];
				entries[i] = l.entries[i];
			}
			_nnz = l._nnz;
		}

		inline sparse_vec(const sparse_vec& l) { copy(l); }

		inline size_t nnz() const { return _nnz; }
		inline size_t size() const { return _nnz; }
		inline size_t alloc() const { return _alloc; }

		sparse_vec(sparse_vec&& l) noexcept {
			indices = l.indices;
			entries = l.entries;
			_nnz = l._nnz;
			_alloc = l._alloc;
			l.indices = nullptr;
			l.entries = nullptr;
			l._nnz = 0;
			l._alloc = 0;
		}

		sparse_vec& operator=(const sparse_vec& l) {
			if (this == &l)
				return *this;

			copy(l);
			return *this;
		}

		sparse_vec& operator=(sparse_vec&& l) noexcept {
			if (this == &l)
				return *this;

			clear();
			indices = l.indices;
			entries = l.entries;
			_nnz = l._nnz;
			_alloc = l._alloc;
			l.indices = nullptr;
			l.entries = nullptr;
			l._nnz = 0;
			l._alloc = 0;
			return *this;
		}

		// this comparison does not clear zero elements / sort the order of indices
		bool operator==(const sparse_vec& l) const {
			if (this == &l)
				return true;
			if (_nnz != l._nnz)
				return false;
			return std::equal(indices, indices + _nnz, l.indices)
				&& std::equal(entries, entries + _nnz, l.entries);
		}

		// this comparison will clear zero elements / sort the order of indices
		bool is_equal_to(const sparse_vec& l) const {
			if (this == &l)
				return true;
			auto this_temp = *this;
			auto other_temp = l;
			this_temp.canonicalize();
			other_temp.canonicalize();
			if (this_temp._nnz != other_temp._nnz)
				return false;
			this_temp.sort_indices();
			other_temp.sort_indices();
			return std::equal(this_temp.indices, this_temp.indices + this_temp._nnz, other_temp.indices)
				&& std::equal(this_temp.entries, this_temp.entries + this_temp._nnz, other_temp.entries);
		}

		inline void push_back(const index_t index, const T& val) {
			if (_nnz + 1 > _alloc)
				reserve((1 + _alloc) * 2); // +1 to avoid _alloc = 0
			indices[_nnz] = index;
			entries[_nnz] = val;
			_nnz++;
		}

		inline void push_back(const index_t index, T&& val) {
			if (_nnz + 1 > _alloc)
				reserve((1 + _alloc) * 2); // +1 to avoid _alloc = 0
			indices[_nnz] = index;
			entries[_nnz] = std::move(val);
			_nnz++;
		}

		void pop_back() {
			if (_nnz != 0)
				_nnz--;
		}

		// take a span of elements
		// sparse_vec.take({start, end}) returns a sparse_vec with elements indexed in [start, end)
		// elements in the resulting sparse_vec are reindexed in [0, end - start)
		sparse_vec<T, index_t> take(const std::pair<index_t, index_t>& span, const bool reserve_nnz = true) const {
			if (span.first < 0 || span.second < 0) {
				report_error("sparse_vec.take: expect non-negative indices.");
				return {};
			}
			if (span.first > span.second) {
				report_error("sparse_vec.take: invalid span.");
				return {};
			}

			sparse_vec<T, index_t> result;
			if (reserve_nnz)
				result.reserve(_nnz);

			for (size_t i = 0; i < _nnz; i++) {
				if (indices[i] >= span.first && indices[i] < span.second) {
					result.push_back(indices[i] - span.first, entries[i]);
				}
			}
			result.compress();
			return result;
		}

		index_t& operator()(const size_t pos) { return indices[pos]; }
		const index_t& operator()(const size_t pos) const { return indices[pos]; }
		T& operator[](const size_t pos) { return entries[pos]; }
		const T& operator[](const size_t pos) const { return entries[pos]; }

		// the entry with this index, or nullptr when there is none
		const T* find(const index_t index, const bool isbinary = true) const {
			if (_nnz == 0)
				return nullptr;
			const index_t* ptr;
			if (isbinary)
				ptr = SparseRREF::binary_search(indices, indices + _nnz, index);
			else
				ptr = std::find(indices, indices + _nnz, index);
			if (ptr == indices + _nnz)
				return nullptr;
			return entries + (ptr - indices);
		}

		T* find(const index_t index, const bool isbinary = true) {
			return const_cast<T*>(std::as_const(*this).find(index, isbinary));
		}

		// conversion functions
		template <typename U = T> requires std::is_integral_v<U> || std::is_same_v<U, int_t>
		operator sparse_vec<rat_t, index_t>() {
			sparse_vec<rat_t, index_t> result;
			result.reserve(_nnz);
			result.resize(_nnz);
			for (size_t i = 0; i < _nnz; i++) {
				result.indices[i] = indices[i];
				result.entries[i] = entries[i];
			}
			return result;
		}

		template <typename U = T> requires std::is_integral_v<U>
		operator sparse_vec<int_t, index_t>() {
			sparse_vec<int_t, index_t> result;
			result.reserve(_nnz);
			result.resize(_nnz);
			for (size_t i = 0; i < _nnz; i++) {
				result.indices[i] = indices[i];
				result.entries[i] = entries[i];
			}
			return result;
		}

		template <typename U = T> requires std::is_same_v<U, rat_t>
		sparse_vec<ulong, index_t> operator%(const nmod_t mod) const {
			sparse_vec<ulong, index_t> result;
			result.reserve(_nnz);
			result.resize(_nnz);
			for (size_t i = 0; i < _nnz; i++) {
				result.indices[i] = indices[i];
				result.entries[i] = entries[i] % mod;
			}
			return result;
		}

		template <typename U = T> requires std::is_same_v<U, rat_t>
		sparse_vec<ulong, index_t> operator%(const ulong p) const {
			nmod_t mod;
			nmod_init(&mod, p);
			return *this % mod;
		}

		void canonicalize() {
			size_t new_nnz = 0;
			for (size_t i = 0; i < _nnz; i++) {
				if (entries[i] != 0) {
					if (new_nnz != i) {
						indices[new_nnz] = indices[i];
						entries[new_nnz] = entries[i];
					}
					new_nnz++;
				}
			}
			_nnz = new_nnz;
		}

		void sort_indices() {
			if (_nnz <= 1 || std::is_sorted(indices, indices + _nnz))
				return;

			auto perm = perm_init(_nnz);
			std::ranges::sort(perm, [&](size_t a, size_t b) {
				return indices[a] < indices[b];
				});

			permute(perm, indices);
			permute(perm, entries);
		}

		void compress() {
			canonicalize();
			sort_indices();
			reserve(_nnz);
		}
	};

	template <Flint::builtin_integral index_t> struct sparse_vec<bool, index_t> {
		index_t* indices = nullptr;
		size_t _nnz = 0;
		size_t _alloc = 0;

		std::span<index_t> index_span() { return std::span<index_t>(indices, _nnz); }
		std::span<const index_t> index_span() const { return std::span<const index_t>(indices, _nnz); }

		sparse_vec() {
			indices = nullptr;
			_nnz = 0;
			_alloc = 0;
		}

		void clear() {
			if (_alloc != 0)
				s_free(indices);
			_alloc = 0;
			_nnz = 0;
		}

		~sparse_vec() { clear(); }

		void reserve(const size_t n) {
			if (n == _alloc)
				return;

			if (n == 0) {
				clear();
				return;
			}

			if (_alloc == 0) {
				indices = s_malloc<index_t>(n);
				_alloc = n;
				return;
			}

			indices = s_realloc(indices, n);
			_alloc = n;
		}

		void resize(const size_t n) { _nnz = n; }

		inline void copy(const sparse_vec& l) {
			if (this == &l)
				return;
			if (_alloc < l._nnz)
				reserve(l._nnz);
			for (size_t i = 0; i < l._nnz; i++) {
				indices[i] = l.indices[i];
			}
			_nnz = l._nnz;
		}

		sparse_vec(const sparse_vec& l) { copy(l); }
		size_t nnz() const { return _nnz; }
		size_t alloc() const { return _alloc; }

		sparse_vec(sparse_vec&& l) noexcept {
			indices = l.indices;
			_nnz = l._nnz;
			_alloc = l._alloc;
			l.indices = nullptr;
			l._nnz = 0;
			l._alloc = 0;
		}

		sparse_vec& operator=(const sparse_vec& l) {
			if (this == &l)
				return *this;

			copy(l);
			return *this;
		}

		sparse_vec& operator=(sparse_vec&& l) noexcept {
			if (this == &l)
				return *this;

			clear();
			indices = l.indices;
			_nnz = l._nnz;
			_alloc = l._alloc;
			l.indices = nullptr;
			l._nnz = 0;
			l._alloc = 0;
			return *this;
		}

		void push_back(const index_t index, const bool val = true) {
			if (_nnz + 1 > _alloc)
				reserve((1 + _alloc) * 2); // +1 to avoid _alloc = 0
			indices[_nnz] = index;
			_nnz++;
		}

		void pop_back() {
			if (_nnz != 0)
				_nnz--;
		}

		index_t& operator()(const size_t pos) { return indices[pos]; }
		const index_t& operator()(const size_t pos) const { return indices[pos]; }
		void zero() { _nnz = 0; }
		void sort_indices() { std::sort(indices, indices + _nnz); }
		void canonicalize() {}
		void compress() { sort_indices(); }
	};

	template <typename T, Flint::builtin_integral index_t = int> struct sparse_mat {
		size_t nrow = 0;
		size_t ncol = 0;
		std::vector<sparse_vec<T, index_t>> rows;

		void init(const size_t r, const size_t c) {
			nrow = r;
			ncol = c;
			rows = std::vector<sparse_vec<T, index_t>>(r);
		}

		sparse_mat() { nrow = 0; ncol = 0; }
		~sparse_mat() {}
		sparse_mat(const size_t r, const size_t c) { init(r, c); }

		sparse_vec<T, index_t>& operator[](const size_t i) { return rows[i]; }
		const sparse_vec<T, index_t>& operator[](const size_t i) const { return rows[i]; }

		T* find(const size_t row, const index_t col, const bool isbinary = true) {
			return rows[row].find(col, isbinary);
		}

		const T* find(const size_t row, const index_t col, const bool isbinary = true) const {
			return rows[row].find(col, isbinary);
		}

		sparse_mat(const sparse_mat& l) {
			init(l.nrow, l.ncol);
			rows = l.rows;
		}

		sparse_mat(sparse_mat&& l) noexcept {
			nrow = l.nrow;
			ncol = l.ncol;
			rows = std::move(l.rows);
			l.nrow = 0;
		}

		sparse_mat& operator=(const sparse_mat& l) {
			if (this == &l)
				return *this;
			nrow = l.nrow;
			ncol = l.ncol;
			rows = l.rows;
			return *this;
		}

		sparse_mat& operator=(sparse_mat&& l) noexcept {
			if (this == &l)
				return *this;
			nrow = l.nrow;
			ncol = l.ncol;
			rows = std::move(l.rows);
			l.nrow = 0;
			return *this;
		}

		void zero() {
			for (size_t i = 0; i < nrow; i++)
				rows[i].zero();
		}

		void clear() {
			for (size_t i = 0; i < nrow; i++) {
				rows[i].clear();
			}
			std::vector<sparse_vec<T, index_t>> tmp;
			rows.swap(tmp); // clear the vector and free memory
			nrow = 0;
			ncol = 0;
		}

		size_t nnz() const {
			size_t n = 0;
			for (size_t i = 0; i < nrow; i++)
				n += rows[i].nnz();
			return n;
		}

		size_t alloc() const {
			size_t n = 0;
			for (size_t i = 0; i < nrow; i++)
				n += rows[i].alloc();
			return n;
		}

		void compress(thread_pool* pool = nullptr) {
			if (pool == nullptr) {
				for (size_t i = 0; i < nrow; i++) {
					rows[i].compress();
				}
			}
			else {
				pool->detach_loop(0, nrow, [&](size_t i) {
					rows[i].compress();
					});
				pool->wait();
			}
		}

		void clear_zero_row() {
			size_t new_nrow = 0;
			for (size_t i = 0; i < nrow; i++) {
				if (rows[i].nnz() != 0) {
					std::swap(rows[new_nrow], rows[i]);
					new_nrow++;
				}
			}
			nrow = new_nrow;
			rows.resize(nrow);
			rows.shrink_to_fit();
		}

		sparse_mat<T, index_t> transpose() const {
			sparse_mat<T, index_t> res(ncol, nrow);
			for (size_t i = 0; i < ncol; i++)
				res[i].zero();

			for (size_t i = 0; i < nrow; i++) {
				for (size_t j = 0; j < rows[i].nnz(); j++) {
					if constexpr (std::is_same_v<T, bool>)
						res[rows[i](j)].push_back(i);
					else
						res[rows[i](j)].push_back(i, rows[i][j]);
				}
			}
			return res;
		}

		void transpose_replace() {
			std::vector<size_t> colnnz(ncol, 0);
			for (size_t i = 0; i < nrow; i++) {
				for (size_t j = 0; j < rows[i].nnz(); j++) {
					colnnz[rows[i](j)]++;
				}
			}

			sparse_mat<T, index_t> res(ncol, nrow);
			for (size_t i = 0; i < nrow; i++) {
				for (size_t j = 0; j < rows[i].nnz(); j++) {
					auto& newrow = res[rows[i](j)];
					newrow.reserve(colnnz[rows[i](j)]);
					if constexpr (std::is_same_v<T, bool>)
						newrow.push_back(i);
					else
						newrow.push_back(i, std::move(rows[i][j]));
				}
				rows[i].clear();
			}
			*this = std::move(res);
		}

		// take a span of rows
		// sparse_mat.take({start, end}) returns a sparse_mat with rows indexed in [start, end)
		sparse_mat<T, index_t> take(const std::pair<size_t, size_t>& span, thread_pool* pool = nullptr) const {
			if (span.second > nrow) {
				report_error("sparse_mat.take: [start, end) out of [0, nrow).");
				return {};
			}
			if (span.first > span.second) {
				report_error("sparse_mat.take: invalid span.");
				return {};
			}

			sparse_mat<T, index_t> res(span.second - span.first, ncol);
			if (pool == nullptr) {
				res.rows.assign(rows.begin() + span.first, rows.begin() + span.second);
			}
			else {
				pool->detach_loop(span.first, span.second, [&](size_t i) {
					res[i - span.first] = rows[i];
					});
				pool->wait();
			}
			return res;
		}

		// take a span of elements
		// sparse_mat.take(levelspec, {start, end}) returns a sparse_mat whose elements have the levelspec-th index in range [start, end)
		// elements in the resulting sparse_mat have their levelspec-th indices reindexed in [0, end - start)
		sparse_mat<T, index_t> take(const size_t levelspec, const std::pair<index_t, index_t>& span, thread_pool* pool = nullptr) const {
			if (span.first < 0 || span.second < 0) {
				report_error("sparse_mat.take: expect non-negative indices.");
				return {};
			}
			if (span.first > span.second) {
				report_error("sparse_mat.take: invalid span.");
				return {};
			}
			if (levelspec > 1) {
				report_error("sparse_mat.take: levelspec must be 0 or 1.");
				return {};
			}

			if (levelspec == 0)
				return take(span, pool);

			// then levelspec == 1
			if (span.second > ncol) {
				report_error("sparse_mat.take: [start, end) out of [0, ncol).");
				return {};
			}
			
			sparse_mat<T, index_t> res(nrow, span.second - span.first);
			if (pool == nullptr) {
				for (size_t i = 0; i < nrow; i++) {
					res.rows[i] = rows[i].take(span);
				}
			}
			else {
				pool->detach_loop(0, nrow, [&](size_t i) {
					res.rows[i] = rows[i].take(span);
					});
				pool->wait();
			}
			return res;
		}

		// if is_binary is true, we assume that the column indices are sorted in each row
		sparse_mat<T, index_t> submat(const std::pair<size_t, size_t>& rowspan,
			const std::pair<size_t, size_t>& colspan, const bool is_binary = true) const {

			if (colspan.second < colspan.first || rowspan.second < rowspan.first) {
				report_error("sparse_mat.submat: invalid span.");
				return sparse_mat<T, index_t>();
			}

			sparse_mat<T, index_t> res(rowspan.second - rowspan.first, colspan.second - colspan.first);

			if (is_binary) {
				for (size_t i = rowspan.first; i < rowspan.second; i++) {
					// binary search for the starting and end position
					auto indptr = rows[i].indices;
					auto it_start = std::lower_bound(indptr, indptr + rows[i].nnz(), colspan.first);
					auto it_end = std::lower_bound(it_start, indptr + rows[i].nnz(), colspan.second);
					res[i - rowspan.first].reserve(it_end - it_start);
					for (auto it = it_start; it != it_end; it++) {
						auto pos = it - indptr;
						res[i - rowspan.first].push_back(indptr[pos] - colspan.first, rows[i].entries[pos]);
					}
				}
			}
			else {
				for (size_t i = rowspan.first; i < rowspan.second; i++) {
					// use more memory but simpler
					res[i - rowspan.first].reserve(rows[i].nnz());
					for (size_t j = 0; j < rows[i].nnz(); j++) {
						auto col_index = rows[i](j);
						if (col_index >= colspan.first && col_index < colspan.second)
							res[i - rowspan.first].push_back(col_index - colspan.first, rows[i][j]);
					}
				}
			}
			return res;
		}

		// append other sparse_mat to this one
		void append(const sparse_mat<T, index_t>& other, thread_pool* pool = nullptr) {
			ncol = std::max(ncol, other.ncol);

			if (pool == nullptr) {
				rows.insert(rows.end(), other.rows.begin(), other.rows.end());
				nrow += other.nrow;
				return;
			}

			rows.reserve(nrow + other.nrow);
			pool->detach_loop(0, other.nrow, [&](size_t i) {
				rows.emplace_back(other.rows[i]);
				});
			pool->wait();
			nrow += other.nrow;
		}

		void append(sparse_mat<T, index_t>&& other) {
			ncol = std::max(ncol, other.ncol);

			rows.insert(rows.end(), std::make_move_iterator(other.rows.begin()), std::make_move_iterator(other.rows.end()));
			nrow += other.nrow;
			other.clear(); // clear the other matrix
		}

		// sort rows by nnz
		void sort_rows_by_nnz() {
			std::ranges::stable_sort(rows, std::less{}, &sparse_vec<T, index_t>::_nnz);
		}

		template <typename U = T> requires std::is_same_v<U, rat_t>
		sparse_mat<ulong> operator%(const nmod_t mod) const {
			sparse_mat<ulong> result(nrow, ncol);
			for (size_t i = 0; i < nrow; i++) {
				result[i] = rows[i] % mod;
			}
			return result;
		}

		template <typename U = T> requires std::is_same_v<U, rat_t>
		ulong height_bits() const {
			ulong max_height = 0;
			for (size_t i = 0; i < nrow; i++) {
				for (size_t j = 0; j < rows[i].nnz(); j++) {
					auto rr = rows[i][j].height_bits();
					if (rr > max_height)
						max_height = rr;
				}
			}

			return max_height;
		}
	};

	// The bookkeeping that every row wise pass of the tensor shares: the entries of a row are counted
	// first, in a vector (one count per row) or in one counter per thread or per block of the row, and
	// the counts then become the rowptr and the total nnz of the result.

	// the total number of entries of a per row count
	inline size_t total_counts(const std::vector<size_t>& counts) {
		return std::accumulate(counts.begin(), counts.end(), (size_t)0);
	}

	// every row's count, summed from the counters it was split over. A counter vector may hold more
	// entries than there are rows (it is indexed by the value of a coordinate, not by the row), so
	// nrow says how many rows to look at
	inline void row_totals(const std::vector<std::vector<size_t>>& block_counts, std::vector<size_t>& counts, const size_t nrow) {
		counts.resize(nrow);
		for (size_t i = 0; i < nrow; i++)
			counts[i] = std::accumulate(block_counts[i].begin(), block_counts[i].end(), (size_t)0);
	}

	// rowptr from the count of every row; rowptr[0] has to be 0 already
	inline void rowptr_from_counts(const std::vector<size_t>& counts, std::vector<size_t>& rowptr) {
		for (size_t i = 0; i < counts.size(); i++)
			rowptr[i + 1] = rowptr[i] + counts[i];
	}

	// the same, with every row's count split over its counters (and the same caveat as row_totals:
	// there may be more counters than there are rows)
	inline void rowptr_from_block_counts(const std::vector<std::vector<size_t>>& block_counts, std::vector<size_t>& rowptr, const size_t nrow) {
		for (size_t i = 0; i < nrow; i++)
			rowptr[i + 1] = rowptr[i] + std::accumulate(block_counts[i].begin(), block_counts[i].end(), (size_t)0);
	}

	// prefix sums in place, when rowptr already holds the count of every row
	inline void rowptr_prefix_sum(std::vector<size_t>& rowptr, const size_t nrow) {
		for (size_t i = 0; i < nrow; i++)
			rowptr[i + 1] += rowptr[i];
	}

	// CSR format for sparse tensor
	template <typename T, typename index_t> struct sparse_tensor_struct {
		size_t rank = 0;
		size_t alloc = 0; // alloc == 0 should imply that colptr and valptr are nullptr
		index_t* colptr = nullptr;
		T* valptr = nullptr;
		std::vector<size_t> dims;
		std::vector<size_t> rowptr;

		using index_v = std::vector<index_t>;
		using index_p = index_t*;
		using const_index_p = const index_t*;
		using interval_t = std::pair<index_t, index_t>;

		// dims and rowptr must not be empty, so that nnz(), check_sorted() and row() are defined on
		// every instance: the empty tensor is dims {0} / rank 1 / rowptr {0,0}, which is safe in
		// both representations (COO nnz() is rowptr[1], CSR nnz() is rowptr[dims[0]])
		void init_empty() {
			dims.assign(1, 0);
			rank = 1;
			rowptr.assign(2, 0);
			alloc = 0;
			colptr = nullptr;
			valptr = nullptr;
		}

		// empty constructor
		sparse_tensor_struct() {
			init_empty();
		}

		// Constructor with dimensions
		// we require that rank >= 2
		void init(const std::vector<size_t>& l, const size_t aoc = 8) {
			dims = l;
			rank = l.size();
			rowptr = std::vector<size_t>(l[0] + 1, 0);
			alloc = aoc;
			colptr = (alloc == 0) ? nullptr : s_malloc<index_t>((rank - 1) * alloc);
			valptr = (alloc == 0) ? nullptr : s_malloc<T>(alloc);
			for (size_t i = 0; i < alloc; i++)
				new (valptr + i) T();
		}

		sparse_tensor_struct(const std::vector<size_t>& l, const size_t aoc = 8) {
			init(l, aoc);
		}

		// Copy constructor
		sparse_tensor_struct(const sparse_tensor_struct& l) {
			init(l.dims, l.alloc);
			std::copy(l.rowptr.begin(), l.rowptr.end(), rowptr.begin());
			std::copy(l.colptr, l.colptr + alloc * (rank - 1), colptr);
			for (size_t i = 0; i < alloc; i++)
				valptr[i] = l.valptr[i];
		}

		// Take over l's buffers (and its rowptr: no O(rank) copy on every move) and leave l empty
		void steal_from(sparse_tensor_struct& l) noexcept {
			dims = std::move(l.dims);
			rowptr = std::move(l.rowptr);
			rank = l.rank;
			alloc = l.alloc;
			colptr = l.colptr;
			valptr = l.valptr;
			l.init_empty();
		}

		// Move constructor
		sparse_tensor_struct(sparse_tensor_struct&& l) noexcept {
			steal_from(l);
		}

		// free the buffers; the capacity drops to 0, dims/rowptr are left alone
		void release() {
			if (alloc == 0)
				return;
			for (size_t i = 0; i < alloc; i++)
				valptr[i].~T();
			s_free(valptr);
			s_free(colptr);
			valptr = nullptr;
			colptr = nullptr;
			alloc = 0;
		}

		// become the canonical empty tensor (O(1): no O(nrow) fill, and dims[0] == 0 makes
		// check_sorted()/sort_indices() skip the rows, so they never touch the null colptr)
		void clear() {
			release();
			init_empty();
		}

		~sparse_tensor_struct() {
			release();
		}

		// IMPORTANT: it only affects the capacity of colptr and valptr and do nothing on 
		// rowptr, dims (and of course nnz), so be careful when using it, it may leads to 
		// alloc < nnz() and cause memory corruption
		void reserve(const size_t size) {
			if (size == alloc)
				return;
			if (size == 0) {
				release(); // capacity only: the shape and the rowptr stay as they are
				return;
			}
			if (alloc == 0) {
				alloc = size;
				colptr = s_malloc<index_t>(size * (rank - 1));
				valptr = s_malloc<T>(size);
				for (size_t i = 0; i < size; i++)
					new (valptr + i) T();
				return;
			}
			colptr = s_realloc<index_t>(colptr, size * (rank - 1));
			if (size > alloc) {
				valptr = s_realloc<T>(valptr, size);
				for (size_t i = alloc; i < size; i++)
					new (valptr + i) T();
			}
			else if (size < alloc) {
				for (size_t i = size; i < alloc; i++)
					valptr[i].~T();
				valptr = s_realloc<T>(valptr, size);
			}
			alloc = size;
		}

		void shrink_to_fit() { reserve(nnz()); }

		// change the dimensions of the tensor
		// it is dangerous, only for internal use
		void change_dims(const std::vector<size_t>& new_dims) {
			dims = new_dims;
			rank = new_dims.size();
			if (alloc != 0)
				colptr = s_realloc<index_t>(colptr, alloc * (rank - 1));
		}

		void zero() {
			if (rank != 0)
				std::fill(rowptr.begin(), rowptr.end(), 0);
		}

		inline size_t nnz() const {
			return rowptr[dims[0]];
		}

		// Copy assignment
		sparse_tensor_struct& operator=(const sparse_tensor_struct& l) {
			if (this == &l)
				return *this;
			auto nz = l.nnz();
			if (alloc == 0) {
				init(l.dims, nz);
				std::copy(l.rowptr.begin(), l.rowptr.end(), rowptr.begin());
				std::copy(l.colptr, l.colptr + nz * (rank - 1), colptr);
				std::copy(l.valptr, l.valptr + nz, valptr);
				return *this;
			}
			const auto old_rank = rank;
			dims = l.dims;
			rank = l.rank;
			rowptr = l.rowptr;
			if (alloc < nz)
				reserve(nz);
			else if (rank != old_rank)
				colptr = s_realloc<index_t>(colptr, alloc * (rank - 1));
			std::copy(l.colptr, l.colptr + nz * (rank - 1), colptr);
			std::copy(l.valptr, l.valptr + nz, valptr);
			return *this;
		}

		// Move assignment
		sparse_tensor_struct& operator=(sparse_tensor_struct&& l) noexcept {
			if (this == &l)
				return *this;
			release(); // everything below is overwritten by steal_from
			steal_from(l);
			return *this;
		}

		std::vector<size_t> row_nums() const {
			return SparseRREF::difference(rowptr);
		}

		// nnz in the i-th row
		size_t row_nnz(const size_t i) const {
			return rowptr[i + 1] - rowptr[i];
		}

		// row index of the i-th entry
		size_t row_index(const size_t i) const {
			auto it = std::upper_bound(rowptr.begin(), rowptr.end(), i);
			return std::distance(rowptr.begin(), it) - 1;
		}

		// remove zero entries, double pointer
		void canonicalize() {
			size_t nnz_now = nnz();
			size_t index = 0;
			std::vector<size_t> newrowptr(dims[0] + 1);
			newrowptr[0] = 0;
			for (size_t i = 0; i < dims[0]; i++) {
				for (size_t j = rowptr[i]; j < rowptr[i + 1]; j++) {
					if (valptr[j] != 0) {
						s_copy(colptr + index * (rank - 1), colptr + j * (rank - 1), rank - 1);
						valptr[index] = valptr[j];
						index++;
					}
				}
				newrowptr[i + 1] = index;
			}
			rowptr = std::move(newrowptr);
		}

		std::pair<index_p, T*> row(const size_t i) {
			return { colptr + rowptr[i] * (rank - 1), valptr + rowptr[i] };
		}

		std::pair<const_index_p, const T*> row(const size_t i) const {
			return { colptr + rowptr[i] * (rank - 1), valptr + rowptr[i] };
		}

		const_index_p entry_lower_bound(const_index_p l) const {
			auto begin = row(l[0]).first;
			auto end = row(l[0] + 1).first;
			if (begin == end)
				return end;
			return SparseRREF::lower_bound(begin, end, l + 1, rank - 1);
		}

		const_index_p entry_lower_bound(const index_v& l) const {
			return entry_lower_bound(l.data());
		}

		index_p entry_lower_bound(const_index_p l) {
			return const_cast<index_p>(std::as_const(*this).entry_lower_bound(l));
		}

		index_p entry_lower_bound(const index_v& l) {
			return const_cast<index_p>(std::as_const(*this).entry_lower_bound(l));
		}

		const_index_p entry_ptr(const_index_p l) const {
			auto ptr = entry_lower_bound(l);
			auto end = row(l[0] + 1).first;
			if (ptr == end || std::equal(ptr, ptr + rank - 1, l + 1))
				return ptr;
			return end;
		}

		const_index_p entry_ptr(const index_v& l) const {
			return entry_ptr(l.data());
		}

		index_p entry_ptr(const_index_p l) {
			return const_cast<index_p>(std::as_const(*this).entry_ptr(l));
		}

		index_p entry_ptr(const index_v& l) {
			return const_cast<index_p>(std::as_const(*this).entry_ptr(l));
		}

		// The one place where an entry is put into the buffer: make room at `index` by moving the whole
		// tail of the buffer up by one entry, write the labels and the value there, and let every row
		// from l[0] on count the new entry
		void insert_at(const index_v& l, const T& val, const size_t index) {
			const size_t nnz = this->nnz();
			if (nnz + 1 > alloc)
				reserve((alloc + 1) * 2);
			for (size_t i = nnz; i > index; i--) {
				auto tmpptr = colptr + (i - 1) * (rank - 1);
				std::copy_backward(tmpptr, tmpptr + (rank - 1), tmpptr + 2 * (rank - 1));
				valptr[i] = valptr[i - 1];
			}
			std::copy(l.begin() + 1, l.begin() + rank, colptr + index * (rank - 1));
			valptr[index] = val;
			for (size_t i = l[0] + 1; i <= dims[0]; i++)
				rowptr[i]++;
		}

		// where an ordered search puts l, and whether that position already holds it (the search never
		// leaves the row of l[0]; a caller that replaces an entry neither of them changes the nnz)
		size_t entry_pos(const index_v& l) const {
			return (entry_lower_bound(l) - colptr) / (rank - 1);
		}

		bool entry_is(const index_v& l, const size_t pos) const {
			if (pos == rowptr[l[0] + 1])
				return false;
			auto ptr = colptr + pos * (rank - 1);
			return std::equal(ptr, ptr + (rank - 1), l.data() + 1);
		}

		// unordered, push back on the end of the row
		void push_back(const index_v& l, const T& val) {
			insert_at(l, val, rowptr[l[0] + 1]);
		}

		// ordered insert: mode = true replaces an entry that is already there, mode = false inserts a
		// second copy of it; insert_add() adds to the entry instead of replacing it
		void insert(const index_v& l, const T& val, const bool mode = true) {
			const size_t pos = entry_pos(l);
			if (mode && entry_is(l, pos)) {
				valptr[pos] = val;
				return;
			}
			insert_at(l, val, pos);
		}

		// ordered add one value
		void insert_add(const index_v& l, const T& val) {
			const size_t pos = entry_pos(l);
			if (entry_is(l, pos)) {
				valptr[pos] += val;
				return;
			}
			insert_at(l, val, pos);
		}

		/**
    	 * @brief Traverse tensor rows in the range `[row_start, row_end)` and apply function `func` to each entry. The function `row_init` is called before processing each row. Non-threaded version.
    	 *
    	 * @tparam F1 The type of the traversal function.
    	 * @tparam F2 The type of the row initialization function.
    	 * @param row_start The starting row index (inclusive).
		 * @param row_end The ending row index (exclusive).
		 * @param func The traversal function to apply to each entry. Takes at least two `size_t` arguments: row index and entry index. May also receive the return value of `row_init` (if not void) as an additional argument.
		 * @param row_init The row initialization function to be called before processing each row. Takes one `size_t` arguments: row index
    	 */
		template <typename F1, typename F2>
		inline void traverse(const size_t row_start, const size_t row_end, F1&& func, F2&& row_init) const {
			using F2_return_type = std::invoke_result_t<F2, size_t>;
			for (size_t i = row_start; i < row_end; i++) {
				if constexpr (std::is_void_v<F2_return_type>) {
					static_assert(std::is_invocable_v<F1, size_t, size_t>, "sparse_tensor_struct.traverse: func must be invocable with (size_t, size_t) when row_init returns void.");
					row_init(i);
					for (size_t j = rowptr[i]; j < rowptr[i + 1]; j++) {
						func(i, j);
					}
				}
				else {
					using F2_lvalue_type = std::conditional_t<std::is_reference_v<F2_return_type>, F2_return_type,std::add_lvalue_reference_t<F2_return_type>>;
					static_assert(std::is_invocable_v<F1, size_t, size_t, F2_lvalue_type>, "sparse_tensor_struct.traverse: func must be invocable with (size_t, size_t, row_init return type) when row_init does not return void.");
					auto row_data = row_init(i);
					for (size_t j = rowptr[i]; j < rowptr[i + 1]; j++) {
						func(i, j, row_data);
					}
				}
			}
		}

		/**
    	 * @brief Traverse tensor rows in the range `[row_start, row_end)` and apply function `func` to each entry using multi-threading. The function `row_init` is called before processing each row, and the function `block_init` is called before processing each block within a row. This will also setup a vector of `entry_blocks_t` for each row to facilitate further parallel operations.
    	 *
    	 * @tparam F1 The type of the traversal function.
    	 * @tparam F2 The type of the row initialization function.
		 * @tparam F3 The type of the block initialization function.
    	 * @param row_start The starting row index (inclusive).
		 * @param row_end The ending row index (exclusive).
		 * @param func The traversal function to apply to each entry. Takes at least two `size_t` arguments: row index and entry index. May also receive the return values of `row_init` and `block_init` (if not void) as additional arguments.
		 * @param row_init The row initialization function to be called before processing each row. Takes two arguments: row index and `entry_blocks_t` for that row.
		 * @param block_init The block initialization function to be called before processing each block within a row. Takes at least two arguments: row index and block index. May also receive the return value of `row_init` (if not void) as an additional argument.
		 * @param pool Pointer to a thread pool for parallel execution. Must not be `nullptr`.
		 * @return A vector of `entry_blocks_t` for each row in the specified range.
    	 */
		template <typename F1, typename F2, typename F3>
		inline std::vector<entry_blocks_t> traverse_setup_blocks(const size_t row_start, const size_t row_end, F1&& func, F2&& row_init, F3&& block_init, thread_pool* pool) const {
			std::vector<entry_blocks_t> row_blocks;
			if (pool == nullptr) {
				report_error("sparse_tensor_struct.traverse_setup_blocks: thread pool is required for setting up blocks.");
				return {};
			}

			using F2_return_type = std::invoke_result_t<F2, size_t, entry_blocks_t>;
			
			auto nthread = pool->get_thread_count();
			for (size_t i = row_start; i < row_end; i++) {
				// separate row elements into blocks
				const entry_blocks_t blks(rowptr[i], rowptr[i + 1], nthread);
				row_blocks.push_back(blks);
				if constexpr (std::is_void_v<F2_return_type>) {
					using F3_return_type = std::invoke_result_t<F3, size_t, size_t>;
					row_init(i, blks);
					pool->detach_sequence(0, blks.nblk, [&](size_t blk) {
						if constexpr (std::is_void_v<F3_return_type>) {
							static_assert(std::is_invocable_v<F1, size_t, size_t, size_t>, "sparse_tensor_struct.traverse_setup_blocks: func must be invocable with (size_t, size_t, size_t) when both row_init and block_init return void.");
							block_init(i, blk);
							for (size_t j = blks.start(blk); j < blks.end(blk); j++) {
								func(i, j, blk);
							}
						}
						else {
							using F3_lvalue_type = std::conditional_t<std::is_reference_v<F3_return_type>, F3_return_type,std::add_lvalue_reference_t<F3_return_type>>;
							static_assert(std::is_invocable_v<F1, size_t, size_t, size_t, F3_lvalue_type>, "sparse_tensor_struct.traverse_setup_blocks: func must be invocable with (size_t, size_t, size_t, block_init return type) when row_init returns void and block_init does not.");
							auto block_data = block_init(i, blk);
							for (size_t j = blks.start(blk); j < blks.end(blk); j++) {
								func(i, j, blk, block_data);
							}
						}
						});
					pool->wait();
				}
				else {
					using F2_lvalue_type = std::conditional_t<std::is_reference_v<F2_return_type>, F2_return_type,std::add_lvalue_reference_t<F2_return_type>>;
					using F3_return_type = std::invoke_result_t<F3, size_t, size_t, F2_lvalue_type>;
					auto row_data = row_init(i, blks);
					pool->detach_sequence(0, blks.nblk, [&](size_t blk) {
						if constexpr (std::is_void_v<std::invoke_result_t<F3, size_t, size_t, F2_lvalue_type>>) {
							static_assert(std::is_invocable_v<F1, size_t, size_t, size_t, std::invoke_result_t<F2, size_t, entry_blocks_t>>, "sparse_tensor_struct.traverse_setup_blocks: func must be invocable with (size_t, size_t, size_t, row_init return type) when block_init returns void and row_init does not.");
							block_init(i, blk, row_data);
							for (size_t j = blks.start(blk); j < blks.end(blk); j++) {
								func(i, j, blk, row_data);
							}
						}
						else {
							using F3_lvalue_type = std::conditional_t<std::is_reference_v<F3_return_type>, F3_return_type,std::add_lvalue_reference_t<F3_return_type>>;
							static_assert(std::is_invocable_v<F1, size_t, size_t, size_t, F2_lvalue_type, F3_lvalue_type>, "sparse_tensor_struct.traverse_setup_blocks: func must be invocable with (size_t, size_t, size_t, row_init return type, block_init return type) when both row_init and block_init do not return void.");
							auto block_data = block_init(i, blk, row_data);
							for (size_t j = blks.start(blk); j < blks.end(blk); j++) {
								func(i, j, blk, row_data, block_data);
							}
						}
						});
					pool->wait();
				}
			}
			
			return row_blocks;
		}

		/**
    	 * @brief Traverse tensor rows in the range `[row_start, row_end)` and apply function `func` to each entry using multi-threading. The function `row_init` is called before processing each row, and the function `block_init` is called before processing each block within a row. This will use the given `std::span<entry_blocks_t>` for multi-threading.
    	 *
    	 * @tparam F1 The type of the traversal function.
    	 * @tparam F2 The type of the row initialization function.
		 * @tparam F3 The type of the block initialization function.
    	 * @param row_start The starting row index (inclusive).
		 * @param row_end The ending row index (exclusive).
		 * @param func The traversal function to apply to each entry. Takes at least two `size_t` arguments: row index and entry index. May also receive the return values of `row_init` and `block_init` (if not void) as additional arguments.
		 * @param row_init The row initialization function to be called before processing each row. Takes two arguments: row index (`size_t`) and `entry_blocks_t` for that row.
		 * @param block_init The block initialization function to be called before processing each block within a row. Takes at least two `size_t` arguments: row index and block index. May also receive the return value of `row_init` (if not void) as an additional argument.
		 * @param row_blocks A `std::span<entry_blocks_t>` containing the blocks for each row to be used.
		 * @param pool Pointer to a thread pool for parallel execution.
    	 */
		template <typename F1, typename F2, typename F3>
		inline void traverse_using_blocks(const size_t row_start, const size_t row_end, F1&& func, F2&& row_init, F3&& block_init, const std::span<entry_blocks_t>& row_blocks, thread_pool* pool) const {
			if (pool == nullptr) {
				report_error("sparse_tensor_struct.traverse_using_blocks: thread pool is required for using blocks.");
				return;
			}

			using F2_return_type = std::invoke_result_t<F2, size_t, entry_blocks_t>;
			
			for (size_t i = row_start; i < row_end; i++) {
				const entry_blocks_t& blks = row_blocks[i - row_start];
				if (blks.nblk == 0)
					continue;
				if constexpr (std::is_void_v<F2_return_type>) {
					using F3_return_type = std::invoke_result_t<F3, size_t, size_t>;
					row_init(i, blks);
					pool->detach_sequence(0, blks.nblk, [&](size_t blk) {
						if constexpr (std::is_void_v<F3_return_type>) {
							static_assert(std::is_invocable_v<F1, size_t, size_t, size_t>, "sparse_tensor_struct.traverse_using_blocks: func must be invocable with (size_t, size_t, size_t) when both row_init and block_init return void.");
							block_init(i, blk);
							for (size_t j = blks.start(blk); j < blks.end(blk); j++) {
								func(i, j, blk);
							}
						}
						else {
							using F3_lvalue_type = std::conditional_t<std::is_reference_v<F3_return_type>, F3_return_type,std::add_lvalue_reference_t<F3_return_type>>;
							static_assert(std::is_invocable_v<F1, size_t, size_t, size_t, F3_lvalue_type>, "sparse_tensor_struct.traverse_using_blocks: func must be invocable with (size_t, size_t, size_t, block_init return type) when row_init returns void and block_init does not.");
							auto block_data = block_init(i, blk);
							for (size_t j = blks.start(blk); j < blks.end(blk); j++) {
								func(i, j, blk, block_data);
							}
						}
						});
					pool->wait();
				}
				else {
					using F2_lvalue_type = std::conditional_t<std::is_reference_v<F2_return_type>, F2_return_type,std::add_lvalue_reference_t<F2_return_type>>;
					using F3_return_type = std::invoke_result_t<F3, size_t, size_t, F2_lvalue_type>;
					auto row_data = row_init(i, blks);
					pool->detach_sequence(0, blks.nblk, [&](size_t blk) {
						if constexpr (std::is_void_v<F3_return_type>) {
							static_assert(std::is_invocable_v<F1, size_t, size_t, size_t, F2_lvalue_type>, "sparse_tensor_struct.traverse_using_blocks: func must be invocable with (size_t, size_t, size_t, row_init return type) when block_init returns void and row_init does not.");
							block_init(i, blk, row_data);
							for (size_t j = blks.start(blk); j < blks.end(blk); j++) {
								func(i, j, blk, row_data);
							}
						}
						else {
							using F3_lvalue_type = std::conditional_t<std::is_reference_v<F3_return_type>, F3_return_type,std::add_lvalue_reference_t<F3_return_type>>;
							static_assert(std::is_invocable_v<F1, size_t, size_t, size_t, F2_lvalue_type, F3_lvalue_type>, "sparse_tensor_struct.traverse_using_blocks: func must be invocable with (size_t, size_t, size_t, row_init return type, block_init return type) when both row_init and block_init do not return void.");
							auto block_data = block_init(i, blk, row_data);
							for (size_t j = blks.start(blk); j < blks.end(blk); j++) {
								func(i, j, blk, row_data, block_data);
							}
						}
						});
					pool->wait();
				}
			}			
		}

		// Row wise map: the kept entries of the rows [row_start, row_end) are counted first and then
		// written into `B`, whose shape the caller has already set. `dest(i)` is the result row of every
		// kept entry of the row i, `keep(i, j)` tells whether the entry (i, j) is copied and
		// `fill(B, i, j, res)` copies it to the running index `res` of its result row: the driver owns the
		// counters (B.rowptr holds them first, so no count vector is allocated) and `fill` only writes and
		// increments.
		//
		// With a pool the entries of every row are split into blocks, so that one very long row is still
		// filled by several threads. A row has a single result row, so all its blocks share it and their
		// counters fit in one flat vector, instead of a matrix with one vector (and one allocation) per
		// row.
		template <typename DestFn, typename KeepFn, typename FillFn>
		void row_map(sparse_tensor_struct<T, index_t>& B, const size_t row_start, const size_t row_end,
			DestFn&& dest, KeepFn&& keep, FillFn&& fill, thread_pool* pool) const {
			const size_t n_out = B.dims[0];
			// the result rowptr holds the counts first; it is zeroed already
			std::vector<size_t>& res_rowptr = B.rowptr;

			if (pool == nullptr) {
				for (size_t i = row_start; i < row_end; i++)
					for (size_t j = rowptr[i]; j < rowptr[i + 1]; j++)
						if (keep(i, j))
							res_rowptr[dest(i) + 1]++;
				rowptr_prefix_sum(res_rowptr, n_out);
				B.reserve(res_rowptr[n_out]);
				// one cursor per result row: its first free index is its rowptr entry
				for (size_t i = row_start; i < row_end; i++) {
					size_t res = res_rowptr[dest(i)];
					for (size_t j = rowptr[i]; j < rowptr[i + 1]; j++)
						if (keep(i, j))
							fill(B, i, j, res);
				}
				return;
			}

			const size_t nrow = row_end - row_start;
			const size_t nthread = pool->get_thread_count();
			// where the counters of every row's blocks start in the flat counter vector
			std::vector<size_t> block_start(nrow + 1, 0);
			for (size_t i = row_start; i < row_end; i++)
				block_start[i - row_start + 1] = block_start[i - row_start] +
					entry_blocks_t(rowptr[i], rowptr[i + 1], nthread).nblk;
			std::vector<size_t> block_offset(block_start[nrow], 0);
			// the counter of a block becomes the first index of that block in the result
			for (size_t i = row_start; i < row_end; i++) {
				const size_t base = block_start[i - row_start];
				const entry_blocks_t blks(rowptr[i], rowptr[i + 1], nthread);
				pool->detach_sequence(0, blks.nblk, [&](size_t blk) {
					size_t n = 0;
					for (size_t j = blks.start(blk); j < blks.end(blk); j++)
						if (keep(i, j))
							n++;
					block_offset[base + blk] = n;
					});
				pool->wait();
			}
			for (size_t i = row_start; i < row_end; i++)
				for (size_t blk = block_start[i - row_start]; blk < block_start[i - row_start + 1]; blk++)
					res_rowptr[dest(i) + 1] += block_offset[blk];
			rowptr_prefix_sum(res_rowptr, n_out);
			B.reserve(res_rowptr[n_out]);
			for (size_t i = row_start; i < row_end; i++) {
				const size_t base = block_start[i - row_start];
				size_t off = res_rowptr[dest(i)];
				for (size_t blk = base; blk < block_start[i - row_start + 1]; blk++) {
					const size_t n = block_offset[blk];
					block_offset[blk] = off;
					off += n;
				}
			}
			for (size_t i = row_start; i < row_end; i++) {
				const size_t base = block_start[i - row_start];
				const entry_blocks_t blks(rowptr[i], rowptr[i + 1], nthread);
				pool->detach_sequence(0, blks.nblk, [&](size_t blk) {
					size_t res = block_offset[base + blk];
					for (size_t j = blks.start(blk); j < blks.end(blk); j++)
						if (keep(i, j))
							fill(B, i, j, res);
					});
				pool->wait();
			}
		}

		// A run of a row map: the `nnz` kept entries of a row that start at the visited entry `k` and
		// all go to the same result row `dest`, written from the index `offset` of that row
		struct run_t {
			size_t dest;
			size_t k;
			size_t nnz;
			size_t offset;
		};

		// Row wise map of a single row, whose entries are visited in the order `perm` gives (nullptr for
		// the natural order). `dest(i, j)` has to be non-decreasing along that order, which a sorted row
		// (and a perm sorted by the same index) gives: then the kept entries of one result row are
		// consecutive, so the serial path writes them in visit order without any counter beyond the
		// result rowptr, and the pool path only has to know where each run of equal result rows starts.
		template <typename DestFn, typename KeepFn, typename FillFn>
		void row_map_run(sparse_tensor_struct<T, index_t>& B, const size_t row, const std::vector<size_t>* perm,
			DestFn&& dest, KeepFn&& keep, FillFn&& fill, thread_pool* pool) const {
			const size_t n_out = B.dims[0];
			const size_t begin = rowptr[row];
			const size_t len = rowptr[row + 1] - begin;
			// the k-th visited entry of the row
			auto entry_at = [&](const size_t k) { return begin + (perm == nullptr ? k : (*perm)[k]); };
			std::vector<size_t>& res_rowptr = B.rowptr;

			if (pool == nullptr) {
				// count: res_rowptr[k + 1] counts the entries whose result row is k
				for (size_t k = 0; k < len; k++) {
					const size_t j = entry_at(k);
					if (keep(row, j))
						res_rowptr[dest(row, j) + 1]++;
				}
				rowptr_prefix_sum(res_rowptr, n_out);
				B.reserve(res_rowptr[n_out]);
				// write: a monotone map puts every kept entry behind the kept entries before it, so its
				// index in the result is the running count, and the result rows come out grouped
				size_t res = 0;
				for (size_t k = 0; k < len; k++) {
					const size_t j = entry_at(k);
					if (keep(row, j))
						fill(B, row, j, res);
				}
				return;
			}

			// pool: one run per group of kept entries with the same result row, so that the writing can be
			// shared out; the counts fall out as the runs are closed
			std::vector<run_t> runs;
			size_t cur_dest = 0;
			size_t cur_nnz = 0;
			bool open = false;
			for (size_t k = 0; k < len; k++) {
				const size_t j = entry_at(k);
				if (!keep(row, j))
					continue;
				const size_t d = dest(row, j);
				if (open && d == cur_dest) {
					cur_nnz++;
					continue;
				}
				if (open) {
					res_rowptr[cur_dest + 1] += cur_nnz;
					runs.back().nnz = cur_nnz;
				}
				runs.push_back(run_t{ d, k, 0, 0 });
				cur_dest = d;
				cur_nnz = 1;
				open = true;
			}
			if (open) {
				res_rowptr[cur_dest + 1] += cur_nnz;
				runs.back().nnz = cur_nnz;
			}
			rowptr_prefix_sum(res_rowptr, n_out);
			B.reserve(res_rowptr[n_out]);
			// the runs are in result row order, so a run starts where its own row is still free
			size_t cursor = 0;
			size_t last_dest = 0;
			bool first = true;
			for (run_t& r : runs) {
				if (first || r.dest != last_dest) {
					first = false;
					last_dest = r.dest;
					cursor = res_rowptr[r.dest];
				}
				r.offset = cursor;
				cursor += r.nnz;
			}
			// write the entries of a run: the run holds exactly `nnz` kept entries from `k` on
			auto write_run = [&](const run_t& r) {
				size_t res = r.offset;
				for (size_t k = r.k, done = 0; done < r.nnz; k++) {
					const size_t j = entry_at(k);
					if (!keep(row, j))
						continue;
					fill(B, row, j, res);
					done++;
				}
				};
			pool->detach_blocks(0, runs.size(), [&](size_t ss, size_t ee) {
				for (size_t r = ss; r < ee; r++)
					write_run(runs[r]);
				});
			pool->wait();
		}


		// take a span of rows
		// sparse_tensor_struct.take({start, end}) returns a sparse_tensor with rows indexed in [start, end)
		sparse_tensor_struct<T, index_t> take(const interval_t& span, thread_pool* pool = nullptr) const {
			if (span.second > dims[0]) {
				report_error("sparse_tensor_struct.take: [start, end) out of [0, dims[0]).");
				return {};
			}
			if (span.first > span.second) {
				report_error("sparse_tensor_struct.take: invalid span.");
				return {};
			}
			if (span.first < 0 || span.second < 0) {
				report_error("sparse_tensor_struct.take: expect non-negative indices.");
				return {};
			}

			std::vector<size_t> res_dims = dims;
			res_dims[0] = span.second - span.first;
			size_t res_nnz = rowptr[span.second] - rowptr[span.first];
			sparse_tensor_struct<T, index_t> B(res_dims, res_nnz);
			const size_t start = rowptr[span.first];
			const size_t end = rowptr[span.second];
			for (size_t i = span.first; i < span.second; i++) {
				B.rowptr[i - span.first + 1] = rowptr[i + 1] - start;
			}

			auto method = [&](size_t ss, size_t ee) {
				std::copy(colptr + ss * (rank - 1), colptr + ee * (rank - 1), B.colptr + (ss - start) * (rank - 1));
				std::copy(valptr + ss, valptr + ee, B.valptr + (ss - start));
				};

			if (pool == nullptr) {
				method(start, end);
			}
			else {
				pool->detach_blocks(start, end, method);
				pool->wait();
			}
			return B;
		}

		// take a span of elements
		// sparse_tensor_struct.take(levelspec, {start, end}) returns a sparse_tensor_struct whose elements have the levelspec-th index in range [start, end)
		// elements in the resulting sparse_tensor_struct have their levelspec-th indices reindexed in [0, end - start)
		sparse_tensor_struct<T, index_t> take(const size_t levelspec, const interval_t& span, thread_pool* pool = nullptr) const {
			if (levelspec == 0)
				return take(span, pool);

			else if (levelspec < rank) {
				if (span.second > dims[levelspec]) {
					report_error("sparse_tensor_struct.take: [start, end) out of [0, dims[levelspec]).");
					return {};
				}
				if (span.first > span.second) {
					report_error("sparse_tensor_struct.take: invalid span.");
					return {};
				}
				if (span.first < 0 || span.second < 0) {
					report_error("sparse_tensor_struct.take: expect non-negative indices.");
					return {};
				}

				std::vector<size_t> res_dims = dims;
				res_dims[levelspec] = span.second - span.first;
				sparse_tensor_struct<T, index_t> B(res_dims);
				// the entries that keep their coordinate at the levelspec-th level, reindexed to the span
				auto dest = [](size_t i) { return i; };
				auto keep = [&](size_t, size_t j) {
					const index_t coord = colptr[j * (rank - 1) + levelspec - 1];
					return coord >= span.first && coord < span.second;
					};
				auto fill = [&](sparse_tensor_struct<T, index_t>& R, size_t, size_t j, size_t& res) {
					auto tmpptr = colptr + j * (rank - 1);
					std::copy(tmpptr, tmpptr + rank - 1, R.colptr + res * (rank - 1));
					R.colptr[res * (rank - 1) + levelspec - 1] -= span.first;
					R.valptr[res] = valptr[j];
					res++;
					};
				row_map(B, 0, dims[0], dest, keep, fill, pool);
				return B;
			}
			else {
				report_error("sparse_tensor_struct.take: levelspec out of rank.");
				return {};
			}
		}

		// extract a (rank-1) tensor by fixing the first index
		sparse_tensor_struct<T, index_t> extract(const index_t index, thread_pool* pool = nullptr) const {
			// Note: require sorted/perm to ensure correctness
			if (index < 0) {
				report_error("sparse_tensor_struct.extract: expect non-negative indices.");
				return {};
			}
			if (index >= static_cast<index_t>(dims[0])) {
				report_error("sparse_tensor_struct.extract: index out of range.");
				return {};
			}

			std::vector<size_t> res_dims(rank - 1);
			for (size_t i = 0; i < rank - 1; i++)
				res_dims[i] = dims[i + 1];
			sparse_tensor_struct<T, index_t> B(res_dims);
			// the result rows are the remaining indices of the row, so the entries have to be visited in
			// increasing order of that index: a perm sorted by the whole index gives it when the row is not
			// sorted already
			const size_t res_nnz = row_nnz(index);
			std::vector<size_t> perm;
			const std::vector<size_t>* order = nullptr;
			if (!check_sorted()) {
				perm = perm_init(res_nnz);
				parallel_sort(perm.begin(), perm.end(), [&](size_t a, size_t b) {
					auto ptra = colptr + (rowptr[index] + a) * (rank - 1);
					auto ptrb = colptr + (rowptr[index] + b) * (rank - 1);
					return lexico_compare(ptra, ptrb, rank - 1) < 0;
					}, pool);
				order = &perm;
			}
			auto dest = [&](size_t, size_t j) { return static_cast<size_t>(colptr[j * (rank - 1)]); };
			auto keep = [](size_t, size_t) { return true; };
			auto fill = [&](sparse_tensor_struct<T, index_t>& R, size_t, size_t j, size_t& res) {
				auto tmpptr = colptr + j * (rank - 1);
				std::copy(tmpptr + 1, tmpptr + (rank - 1), R.colptr + res * (rank - 2));
				R.valptr[res] = valptr[j];
				res++;
				};
			row_map_run(B, index, order, dest, keep, fill, pool);
			return B;
		}

		// extract a (rank-1) tensor by fixing the levelspec-th index
		sparse_tensor_struct<T, index_t> extract(const size_t levelspec, const index_t index, thread_pool* pool = nullptr) const {
			// Note: require sorted/perm to ensure correctness if levelspec == 0
			if (levelspec == 0)
				return extract(index, pool);
			if (levelspec >= rank) {
				report_error("sparse_tensor_struct.extract: levelspec out of rank.");
				return {};
			}
			if (index < 0) {
				report_error("sparse_tensor_struct.extract: expect non-negative indices.");
				return {};
			}
			if (index >= static_cast<index_t>(dims[levelspec])) {
				report_error("sparse_tensor_struct.extract: index out of range.");
				return {};
			}

			std::vector<size_t> res_dims(rank - 1);
			for (size_t i = 0; i < rank - 1; i++) {
				if (i < levelspec)
					res_dims[i] = dims[i];
				else
					res_dims[i] = dims[i + 1];
			}
			sparse_tensor_struct<T, index_t> B(res_dims);
			auto dest = [](size_t i) { return i; };
			auto keep = [&](size_t, size_t j) { return colptr[j * (rank - 1) + levelspec - 1] == index; };
			auto fill = [&](sparse_tensor_struct<T, index_t>& R, size_t, size_t j, size_t& res) {
				auto tmpptr = colptr + j * (rank - 1);
				std::copy(tmpptr, tmpptr + levelspec - 1, R.colptr + res * (rank - 2));
				std::copy(tmpptr + levelspec, tmpptr + (rank - 1), R.colptr + res * (rank - 2) + levelspec - 1);
				R.valptr[res] = valptr[j];
				res++;
				};
			row_map(B, 0, dims[0], dest, keep, fill, pool);
			return B;
		}

		/* extract certain part from tensor, similar to Part in Mathematica 
		   part specification for each level:
		   - `index_t`: extract that index
		   - `std::pair<index_t, index_t>`: take the span `[first, second)`
		   Note: the result must be itself a sparse tensor_struct, i.e., has rank >= 2.
		*/
		sparse_tensor_struct<T, index_t> part(const std::vector<std::variant<index_t, interval_t>>& partspec, thread_pool* pool = nullptr) const {
			// Note: require sorted/perm to ensure correctness if partspec[0] is index_t
			if (partspec.size() != rank) {
				report_error("sparse_tensor_struct.part: partspec size not equal to rank.");
				return {};
			}

			// process partspec to get result dimensions
			std::vector<size_t> res_dims;
			std::vector<uint8_t> is_span(rank, 0); // record the partspec type at each level
			for (size_t i = 0; i < rank; i++) {
				if (std::holds_alternative<index_t>(partspec[i])) {
					index_t idx = std::get<index_t>(partspec[i]);
					if (idx < 0 || idx >= static_cast<index_t>(dims[i])) {
						report_error("sparse_tensor_struct.part: index out of range.");
						return {};
					}
					is_span[i] = 0;
				}
				else {
					auto& span = std::get<interval_t>(partspec[i]);
					if (span.first < 0 || span.second < 0) {
						report_error("sparse_tensor_struct.part: expect non-negative indices.");
						return {};
					}
					if (span.first > span.second) {
						report_error("sparse_tensor_struct.part: invalid span.");
						return {};
					}
					if (span.second > static_cast<index_t>(dims[i])) {
						report_error("sparse_tensor_struct.part: [start, end) out of [0, dims[i]).");
						return {};
					}
					if (span.first == 0 && span.second == dims[i])
						is_span[i] = 2; // full span
					else
						is_span[i] = 1; // partial span
					res_dims.push_back(span.second - span.first);
				}
			}
			// check result rank
			if (res_dims.size() < 2) {
				report_error("sparse_tensor_struct.part: result tensor rank < 2.");
				return {};
			}

			// Now introduce some helper variables and functions

			// levels which need to check (index or partial span), element: [level, type]
			std::vector<std::pair<size_t, uint8_t>> coord_check_levels;
			for (size_t i = 1; i < rank; i++) {
				if (is_span[i] != 2) {
					coord_check_levels.push_back(std::make_pair(i, is_span[i]));
				}
			}
			// check if an coordinate is in the partspec
			auto check_coordinate = [&](const_index_p coord) {
				for (const auto& [level, type] : coord_check_levels) {
					if (type == 0) {
						const index_t idx = std::get<index_t>(partspec[level]);
						if (coord[level - 1] != idx)
							return false;
					}
					else if (type == 1) {
						const auto& span = std::get<interval_t>(partspec[level]);
						if (coord[level - 1] < span.first || coord[level - 1] >= span.second)
							return false;
					}
				}
				return true;
			};

			// levels which need to copy (span), element: [[old level, new level - 1], reindex offset] (note that new level - 1 == its position in new_coord)
			std::vector<std::pair<std::pair<size_t, size_t>, index_t>> coord_copy_levels;
			for (size_t i = 1; i < rank; i++) {
				if (is_span[i] != 0) {
					coord_copy_levels.push_back(std::make_pair(std::make_pair(i, coord_copy_levels.size()), std::get<interval_t>(partspec[i]).first));
				}
			}
			// used when is_span[0] == 0, old colptr[new_row_level] becomes the row level (therefore chopped in new colptr)
			const size_t new_row_level = coord_copy_levels[0].first.first - 1;
			const index_t new_row_reindex = coord_copy_levels[0].second;
			// when is_span[0] == 0, old colptr[new_row_level] should not be copied to new colptr
			if (is_span[0] == 0) {
				coord_copy_levels.erase(coord_copy_levels.begin());
				// re-index new levels
				for (auto& [old_new_level, offset] : coord_copy_levels) {
					old_new_level.second -= 1;
				}
			}
			// copy colptr from old_coord to new_coord
			auto copy_colptr = [&](const_index_p old_coord, index_p new_coord) {
				for (const auto& [old_new_level, reindex] : coord_copy_levels) {
					auto& [old_level, new_pos] = old_new_level;
					new_coord[new_pos] = old_coord[old_level - 1] - reindex;
				}
			};

			// interval representation of rowspec
			interval_t row_span;
			if (is_span[0] == 0) {
				index_t idx = std::get<index_t>(partspec[0]);
				row_span = std::make_pair(idx, idx + 1);
			}
			else {
				row_span = std::get<interval_t>(partspec[0]);
			}
			// rowspec index, only used when is_span[0] == 0
			const size_t rowspec_idx = (is_span[0] == 0) ? std::get<index_t>(partspec[0]) : row_span.first;

			const bool sorted = check_sorted();
			sparse_tensor_struct<T, index_t> B(res_dims);

			// the result row of an entry: the rowspec fixes it when it is a span, and the first copied
			// level gives it when the rowspec is an index (then every entry read comes from one row)
			auto dest_span = [&](size_t i) { return i - row_span.first; };
			auto dest_index = [&](size_t, size_t j) {
				return static_cast<size_t>(colptr[j * (rank - 1) + new_row_level] - new_row_reindex);
				};
			auto keep = [&](size_t, size_t j) { return check_coordinate(colptr + j * (rank - 1)); };
			auto fill = [&](sparse_tensor_struct<T, index_t>& R, size_t, size_t j, size_t& res) {
				auto oldptr = colptr + j * (rank - 1);
				auto newptr = R.colptr + res * (R.rank - 1);
				copy_colptr(oldptr, newptr);
				R.valptr[res] = valptr[j];
				res++;
				};

			// `Sorted` is false only when the rowspec is an index and the tensor is not sorted: then the
			// entries of the one row we read have to be permuted into result row order first
			auto part_impl = [&]<bool RowspecIsSpan, bool Sorted>() {
				if constexpr (RowspecIsSpan) {
					row_map(B, row_span.first, row_span.second, dest_span, keep, fill, pool);
				}
				else {
					std::vector<size_t> perm;
					const std::vector<size_t>* order = nullptr;
					if constexpr (!Sorted) {
						perm = perm_init(rowptr[rowspec_idx + 1] - rowptr[rowspec_idx]);
						parallel_sort(perm.begin(), perm.end(), [&](size_t a, size_t b) {
							auto ptra = colptr + (rowptr[rowspec_idx] + a) * (rank - 1);
							auto ptrb = colptr + (rowptr[rowspec_idx] + b) * (rank - 1);
							return lexico_compare(ptra, ptrb, rank - 1) < 0;
							}, pool);
						order = &perm;
					}
					row_map_run(B, rowspec_idx, order, dest_index, keep, fill, pool);
				}
				};
			if (is_span[0] != 0) {
				if (sorted)
					part_impl.template operator()<true, true>();
				else
					part_impl.template operator()<true, false>();
			}
			else if (sorted)
				part_impl.template operator()<false, true>();
			else
				part_impl.template operator()<false, false>();
			return B;
		}

		sparse_tensor_struct<T, index_t> transpose(const std::vector<size_t>& perm, thread_pool* pool = nullptr, const bool sort_ind = true) const {
			std::vector<size_t> l(rank);
			std::vector<size_t> lperm(rank);
			std::vector<index_t> old_coord(rank);
			std::vector<index_t> new_coord(rank);
			for (size_t i = 0; i < rank; i++)
				lperm[i] = dims[perm[i]];
			sparse_tensor_struct<T, index_t> B(lperm, nnz());
			for (size_t i = 0; i < dims[0]; i++) {
				for (size_t j = rowptr[i]; j < rowptr[i + 1]; j++) {
					old_coord[0] = i;
					auto tmpptr = colptr + j * (rank - 1);
					for (size_t k = 1; k < rank; k++)
						old_coord[k] = tmpptr[k - 1];
					for (size_t k = 0; k < rank; k++)
						new_coord[k] = old_coord[perm[k]];
					B.push_back(new_coord, valptr[j]);
				}
			}
			if (sort_ind)
				B.sort_indices(pool);
			return B;
		}

		// the entries of every row have to be non-decreasing. Comparing adjacent entries directly
		// avoids the perm vector and the parallel adjacent_find that the previous version paid for
		// every row
		bool check_sorted() const {
			for (size_t i = 0; i < dims[0]; i++)
				for (size_t j = rowptr[i] + 1; j < rowptr[i + 1]; j++)
					if (lexico_compare(colptr + (j - 1) * (rank - 1), colptr + j * (rank - 1), rank - 1) > 0)
						return false;
			return true;
		}

		// multithread version use more memory and it will compress the tensor
		void sort_indices(thread_pool* pool = nullptr) {
			auto sort_row = [&](const size_t row, std::vector<size_t>& perm) {
				const size_t rownnz = rowptr[row + 1] - rowptr[row];
				perm.resize(rownnz);
				for (size_t j = 0; j < rownnz; j++)
					perm[j] = j;
				auto less = [&](size_t a, size_t b) {
					auto ptra = colptr + (rowptr[row] + a) * (rank - 1);
					auto ptrb = colptr + (rowptr[row] + b) * (rank - 1);
					return lexico_compare(ptra, ptrb, rank - 1) < 0;
					};
				parallel_sort(perm.begin(), perm.end(), less, pool);
				};

			if (pool == nullptr) {
				std::vector<size_t> perm;
				for (size_t i = 0; i < dims[0]; i++) {
					sort_row(i, perm);
					permute(perm, colptr + rowptr[i] * (rank - 1), rank - 1);
					permute(perm, valptr + rowptr[i]);
				}
				return;
			}

			auto nz = nnz();
			if (nz == 0) {
				release(); // nothing to sort: keep the shape, just drop the unused capacity
				return;
			}
			auto n_colptr = s_malloc<index_t>(nz * (rank - 1));
			auto n_valptr = s_malloc<T>(nz);
			for (size_t i = 0; i < nz; i++)
				new (n_valptr + i) T();

			auto loop_rows = [&](size_t ss, size_t ee) {
				std::vector<size_t> perm;
				for (size_t i = ss; i < ee; i++) {
					sort_row(i, perm);
					for (size_t j = 0; j < perm.size(); j++) {
						auto oldptr = colptr + (rowptr[i] + perm[j]) * (rank - 1);
						auto newptr = n_colptr + (rowptr[i] + j) * (rank - 1);
						std::copy(oldptr, oldptr + (rank - 1), newptr);
						n_valptr[rowptr[i] + j] = valptr[rowptr[i] + perm[j]];
					}
				}
				};

			if (dims[0] < pool->get_thread_count()) {
				// few rows: each row gets the pool (and a long one is sorted in parallel), instead of
				// the rows being shared out (which would leave a single row to one thread)
				loop_rows(0, dims[0]);
			}
			else {
				// one dispatch for all the rows, one perm vector per block instead of one per row
				pool->detach_blocks(0, dims[0], [&](size_t ss, size_t ee) { loop_rows(ss, ee); });
				pool->wait();
			}
			for (size_t i = 0; i < alloc; i++)
				valptr[i].~T();
			s_free(colptr);
			s_free(valptr);
			colptr = n_colptr;
			valptr = n_valptr;
			alloc = nz;
		}
	};

	// define the default sparse tensor
	template <typename T, typename index_t = int, SPARSE_TYPE Type = SPARSE_COO> struct sparse_tensor;

	template <typename T, typename index_t> struct sparse_tensor<T, index_t, SPARSE_CSR> {
		sparse_tensor_struct<T, index_t> data;

		using index_v = std::vector<index_t>;
		using index_p = index_t*;
		using const_index_p = const index_t*;

		void clear() { data.clear(); }

		// the empty state is a valid 0 x 0 matrix, so that nnz(), dim(), check_sorted() and the
		// rowptr bookkeeping stay well defined on a tensor returned by a rejected call
		sparse_tensor() {
			data.rank = 2;
			data.dims = { 0, 0 };
			data.rowptr = { 0 };
		}
		~sparse_tensor() {}
		sparse_tensor(const std::vector<size_t>& l, size_t aoc = 8) : data(l, aoc) {}
		sparse_tensor(const sparse_tensor& l) : data(l.data) {}
		sparse_tensor(sparse_tensor&& l) noexcept : data(std::move(l.data)) {}
		sparse_tensor& operator=(const sparse_tensor& l) { data = l.data; return *this; }
		sparse_tensor& operator=(sparse_tensor&& l) noexcept { data = std::move(l.data); return *this; }

		size_t alloc() const { return data.alloc; }
		size_t rank() const { return data.rank; }
		size_t nnz() const { return data.rowptr[data.dims[0]]; }
		auto& rowptr() { return data.rowptr; }
		const auto& rowptr() const { return data.rowptr; }
		index_p colptr() { return data.colptr; }
		const_index_p colptr() const { return data.colptr; }
		T* valptr() { return data.valptr; }
		const T* valptr() const { return data.valptr; }
		auto& dims() { return data.dims; }
		const auto& dims() const { return data.dims; }
		size_t dim(const size_t i) const { return data.dims[i]; }
		index_p index(const size_t i) { return data.colptr + i * (rank() - 1); }
		const_index_p index(const size_t i) const { return data.colptr + i * (rank() - 1); }
		T& val(const size_t i) { return data.valptr[i]; }
		const T& val(const size_t i) const { return data.valptr[i]; }
		bool check_sorted() const { return data.check_sorted(); }
		void zero() { data.zero(); }
		void insert(const index_v& l, const T& val, bool mode = true) { data.insert(l, val, mode); }
		void push_back(const index_v& l, const T& val) { data.push_back(l, val); }
		void canonicalize() { data.canonicalize(); }
		void sort_indices(thread_pool* pool = nullptr) { data.sort_indices(pool); }
		void reserve(const size_t size) { data.reserve(size); }
		void shrink_to_fit() { data.shrink_to_fit(); }
		sparse_tensor transpose(const std::vector<size_t>& perm, thread_pool* pool = nullptr, const bool sort_ind = true) const {
			sparse_tensor B;
			B.data = data.transpose(perm, pool, sort_ind);
			return B;
		}

		void change_dims(const std::vector<size_t>& new_dims) { data.change_dims(new_dims); }

		// index vector (row index included at first) of the i-th entry
		index_v index_vector(const size_t i) const {
			index_v result(rank());
			result[0] = static_cast<index_t>(data.row_index(i));
			for (size_t j = 1; j < rank(); j++)
				result[j] = index(i)[j - 1];
			return result;
		}

		// take a span of elements
		// sparse_tensor.take(levelspec, {start, end}) returns a sparse_tensor whose elements have the levelspec-th index in range [start, end)
		// elements in the resulting sparse_tensor have their levelspec-th indices reindexed in [0, end - start)
		sparse_tensor take(const size_t levelspec, const std::pair<size_t, size_t>& span, thread_pool* pool = nullptr) const {
			sparse_tensor B;
			B.data = data.take(levelspec, span, pool);
			return B;
		}

		// extract a (rank-1) tensor by fixing the levelspec-th index
		sparse_tensor extract(const size_t levelspec, const size_t index, thread_pool* pool = nullptr) const {
			// Note: require sorted/perm to ensure correctness if levelspec == 0
			sparse_tensor B;
			B.data = data.extract(levelspec, index, pool);
			return B;
		}

		/* extract certain part from tensor, similar to Part in Mathematica 
		   part specification for each level:
		   - `index_t`: extract that index
		   - `std::vector<std::pair<index_t, index_t>>`: take the union of spans `[first, second)`
		   Note: the result must be itself a sparse_tensor, i.e., has rank >= 2.
		*/
		sparse_tensor part(const std::vector<std::variant<index_t, std::pair<index_t, index_t>>>& partspec, thread_pool* pool = nullptr) const {
			// Note: require sorted/perm to ensure correctness if partspec[0] is index_t
			sparse_tensor B;
			B.data = data.part(partspec, pool);
			return B;
		}

		void convert_from_COO(const sparse_tensor<T, index_t, SPARSE_COO>& l, thread_pool* pool = nullptr) {
			// a rank 0 tensor has no CSR form: there is no row index to build the rowptr from, and the
			// data.dims[0] reads below would be out of bounds. The empty matrix is all that CSR can
			// express here, so report the lost entries and return it
			if (l.rank() == 0) {
				if (l.nnz() != 0)
					report_error("sparse_tensor: a rank 0 tensor cannot be converted to CSR");
				data.clear();
				data.init({ 0, 0 });
				return;
			}
			// Note: require sorted/perm to ensure correctness
			const auto old_rank = data.rank;
			data.dims = l.data.dims;
			data.rank = l.data.rank;
			auto nnz = l.nnz();
			auto newrank = data.rank - 1;
			for (size_t i = 0; i < newrank; i++)
				data.dims[i] = data.dims[i + 1];
			data.dims.resize(newrank);
			data.rank = newrank;
			if (alloc() < nnz)
				reserve(nnz);
			else if (old_rank != l.rank() && alloc() != 0)
				data.colptr = s_realloc<index_t>(data.colptr, alloc() * (l.rank() - 1));
			std::copy(l.data.valptr, l.data.valptr + nnz, data.valptr);
			// then recompute the rowptr and colptr
			// first compute nnz for each row
			data.rowptr.assign(data.dims[0] + 1, 0);
			bool l_sorted = l.check_sorted();
			if (l_sorted || pool == nullptr) {
				for (size_t i = 0; i < nnz; i++) {
					auto oldptr = l.data.colptr + i * newrank;
					auto nowptr = data.colptr + i * (newrank - 1);
					data.rowptr[oldptr[0] + 1]++;
					std::copy(oldptr + 1, oldptr + newrank, nowptr);
				}
				rowptr_prefix_sum(data.rowptr, data.dims[0]);
			}
			if (!l_sorted) {
				std::vector<size_t> perm = perm_init(nnz);
				auto by_index = [&](size_t a, size_t b) {
					auto ptra = l.data.colptr + a * newrank;
					auto ptrb = l.data.colptr + b * newrank;
					return lexico_compare(ptra, ptrb, newrank) < 0;
					};
				parallel_sort(perm.begin(), perm.end(), by_index, pool);
				if (pool == nullptr) {
					permute(perm, data.colptr, newrank - 1);
					permute(perm, data.valptr);
				}
				else {
					auto nthread = pool->get_thread_count();
					std::vector<std::vector<size_t>> row_block_nnz(data.dims[0] + 1, std::vector<size_t>(nthread, 0));
					pool->detach_blocks(0, nnz, [&](size_t ss, size_t ee) {
						for (size_t j = ss; j < ee; j++) {
							auto oldptr = l.data.colptr + perm[j] * newrank;
							auto newptr = data.colptr + j * (newrank - 1);
							auto id = SparseRREF::thread_id();
							row_block_nnz[oldptr[0]][id]++;
							std::copy(oldptr + 1, oldptr + newrank, newptr);
							data.valptr[j] = l.data.valptr[perm[j]];
						}
					});
					pool->wait();
					rowptr_from_block_counts(row_block_nnz, data.rowptr, data.dims[0]);
				}
			}
		}

		void move_from_COO(sparse_tensor<T, index_t, SPARSE_COO>&& l, thread_pool* pool = nullptr) noexcept {
			// a rank 0 tensor has no CSR form: there is no row index to build the rowptr from, and the
			// data.dims[0] reads below would be out of bounds. The empty matrix is all that CSR can
			// express here, so report the lost entries and return it
			if (l.rank() == 0) {
				if (l.nnz() != 0)
					report_error("sparse_tensor: a rank 0 tensor cannot be converted to CSR");
				data.clear();
				data.init({ 0, 0 });
				return;
			}
			// Note: require sorted/perm to ensure correctness
			// use move to avoid memory problems, then no need to mention l
			data = std::move(l.data);
			// ensure sorted
			if (!data.check_sorted())
				data.sort_indices(pool);
			// remove the first dimension
			auto newrank = data.rank - 1;
			for (size_t i = 0; i < newrank; i++)
				data.dims[i] = data.dims[i + 1];
			data.dims.resize(newrank);
			// then recompute the rowptr and colptr
			// first compute nnz for each row
			std::vector<size_t> rowptr(data.dims[0] + 1, 0);
			auto nnz = data.rowptr[1];
			for (size_t i = 0; i < nnz; i++) {
				auto oldptr = data.colptr + i * newrank;
				auto nowptr = data.colptr + i * (newrank - 1);
				rowptr[oldptr[0] + 1]++;
				for (size_t j = 0; j < newrank - 1; j++)
					nowptr[j] = oldptr[j + 1];
			}
			rowptr_prefix_sum(rowptr, data.dims[0]);
			data.rowptr = std::move(rowptr);
			data.rank = newrank;
			// colptr only holds nnz * (rank - 1) indices now, so alloc must not claim more than that:
			// a later copy reads exactly alloc * (rank - 1) of them
			if (data.alloc == nnz && nnz != 0)
				data.colptr = s_realloc<index_t>(data.colptr, nnz * (data.rank - 1));
			else
				data.reserve(nnz);
		}

		// constructor from COO
		sparse_tensor(const sparse_tensor<T, index_t, SPARSE_COO>& l, thread_pool* pool = nullptr) { convert_from_COO(l, pool); }
		sparse_tensor& operator=(const sparse_tensor<T, index_t, SPARSE_COO>& l) {
			// Note: require sorted/perm to ensure correctness
			convert_from_COO(l);
			return *this;
		}

		// suppose that COO tensor is sorted
		sparse_tensor(sparse_tensor<T, index_t, SPARSE_COO>&& l, thread_pool* pool = nullptr) noexcept {
			// Note: require sorted/perm to ensure correctness
			move_from_COO(std::move(l), pool);
		}

		sparse_tensor& operator=(sparse_tensor<T, index_t, SPARSE_COO>&& l) noexcept {
			// Note: require sorted/perm to ensure correctness
			move_from_COO(std::move(l));
			return *this;
		}

		// convert a sparse matrix to a sparse tensor
		void convert_from_sparse_mat(const sparse_mat<T, index_t>& mat, thread_pool* pool = nullptr, const bool sort_ind = true) {
			data.dims = { mat.nrow, mat.ncol };
			data.rank = 2;

			auto nnz = mat.nnz();
			if (alloc() < nnz)
				reserve(nnz);

			// compute the rowptr
			data.rowptr.resize(mat.nrow + 1, 0);
			for (size_t i = 0; i < mat.nrow; i++) {
				data.rowptr[i + 1] = data.rowptr[i] + mat[i].nnz();
			}
			// copy the values and column indices
			// one row body, run by the pool or in a plain loop
			auto copy_row = [&](const size_t i) {
				std::copy(mat[i].indices, mat[i].indices + mat[i].nnz(), data.colptr + data.rowptr[i]);
				std::copy(mat[i].entries, mat[i].entries + mat[i].nnz(), data.valptr + data.rowptr[i]);
				};
			if (pool == nullptr) {
				for (size_t i = 0; i < mat.nrow; i++)
					copy_row(i);
			}
			else {
				pool->detach_loop(0, mat.nrow, copy_row);
				pool->wait();
			}
			if (sort_ind && !check_sorted())
				sort_indices(pool);
		}

		void move_from_sparse_mat(sparse_mat<T, index_t>&& mat, thread_pool* pool = nullptr, const bool sort_ind = true) {
			data.clear();
			data.init({ mat.nrow, mat.ncol }, mat.nnz());

			for (size_t i = 0; i < mat.nrow; i++)
				data.rowptr[i + 1] = data.rowptr[i] + mat[i].nnz();

			auto move_row = [&](size_t i) {
				if (mat[i].nnz() != 0) {
					if (sort_ind)
						mat[i].sort_indices(); // sort each source row in place
					std::copy(mat[i].indices, mat[i].indices + mat[i].nnz(), data.colptr + data.rowptr[i]);
					std::move(mat[i].entries, mat[i].entries + mat[i].nnz(), data.valptr + data.rowptr[i]);
				}
				mat[i].clear();
				};
			if (pool == nullptr) {
				for (size_t i = 0; i < mat.nrow; i++)
					move_row(i);
			}
			else {
				pool->detach_loop(0, mat.nrow, move_row);
				pool->wait();
			}
			mat.clear();
		}

		sparse_mat<T, index_t> to_sparse_mat(thread_pool* pool = nullptr) const & {
			if (rank() != 2) {
				report_error("sparse_tensor.to_sparse_mat: rank must be 2");
				return sparse_mat<T, index_t>();
			}

			sparse_mat<T, index_t> mat(data.dims[0], data.dims[1]);
			// one row body, run by the pool or in a plain loop
			auto copy_row = [&](const size_t i) {
				const size_t nz = data.rowptr[i + 1] - data.rowptr[i];
				mat[i].reserve(nz);
				mat[i].resize(nz);
				std::copy(data.colptr + data.rowptr[i], data.colptr + data.rowptr[i + 1], mat[i].indices);
				std::copy(data.valptr + data.rowptr[i], data.valptr + data.rowptr[i + 1], mat[i].entries);
				};
			if (pool == nullptr) {
				for (size_t i = 0; i < data.dims[0]; i++)
					copy_row(i);
			}
			else {
				pool->detach_loop(0, data.dims[0], copy_row);
				pool->wait();
			}
			return mat;
		}

		sparse_mat<T, index_t> to_sparse_mat(thread_pool* pool = nullptr) && {
			if (rank() != 2) {
				std::cerr << "sparse_tensor.to_sparse_mat: rank must be 2" << std::endl;
				return sparse_mat<T, index_t>();
			}

			sparse_mat<T, index_t> mat(data.dims[0], data.dims[1]);
			auto move_row = [&](size_t i) {
				auto nz = data.rowptr[i + 1] - data.rowptr[i];
				if (nz == 0)
					return;
				mat[i].reserve(nz);
				mat[i].resize(nz);
				std::copy(data.colptr + data.rowptr[i], data.colptr + data.rowptr[i + 1], mat[i].indices);
				std::move(data.valptr + data.rowptr[i], data.valptr + data.rowptr[i + 1], mat[i].entries);
				};
			if (pool == nullptr) {
				for (size_t i = 0; i < data.dims[0]; i++)
					move_row(i);
			}
			else {
				pool->detach_loop(0, data.dims[0], move_row);
				pool->wait();
			}
			clear();
			return mat;
		}

		sparse_tensor(const sparse_mat<T, index_t>& mat, thread_pool* pool = nullptr) {
			convert_from_sparse_mat(mat, pool);
		}

		sparse_tensor& operator=(const sparse_mat<T, index_t>& mat) {
			convert_from_sparse_mat(mat);
			return *this;
		}

		sparse_tensor(sparse_mat<T, index_t>&& mat, thread_pool* pool = nullptr) {
			move_from_sparse_mat(std::move(mat), pool);
		}

		sparse_tensor& operator=(sparse_mat<T, index_t>&& mat) {
			move_from_sparse_mat(std::move(mat));
			return *this;
		}

		// only for test
		void print_test() const {
			for (size_t i = 0; i < data.dims[0]; i++) {
				for (size_t j = data.rowptr[i]; j < data.rowptr[i + 1]; j++) {
					std::cout << i << " ";
					for (size_t k = 0; k < data.rank - 1; k++)
						std::cout << (size_t)data.colptr[j * (data.rank - 1) + k] << " ";
					std::cout << " : " << data.valptr[j] << std::endl;
				}
			}
		}
	};

	template <typename index_t, typename T> struct sparse_tensor<T, index_t, SPARSE_COO> {
		sparse_tensor_struct<T, index_t> data;

		using index_v = std::vector<index_t>;
		using index_p = index_t*;
		using const_index_p = const index_t*;

		template <typename S, typename U = S> requires std::convertible_to<U, S>
		std::vector<S> prepend_num(const std::vector<S>& l, U num = 0) const {
			std::vector<S> lp;
			lp.reserve(l.size() + 1);
			lp.push_back(static_cast<S>(num));
			lp.insert(lp.end(), l.begin(), l.end());
			return lp;
		}

		void clear() { data.clear(); }
		void init(const std::vector<size_t>& l, size_t aoc = 8) {
			data.init(prepend_num(l, (size_t)1), aoc);
		}

		// the empty state is a valid rank 0 tensor, so that nnz(), rank(), dims() and
		// check_sorted() stay well defined on a tensor returned by a failed call
		sparse_tensor() {
			data.rank = 1;
			data.dims = { 1 };
			data.rowptr = { 0, 0 };
		}
		~sparse_tensor() {}
		sparse_tensor(const std::vector<size_t>& l, size_t aoc = 8) : data(prepend_num(l, (size_t)1), aoc) {}
		sparse_tensor(const sparse_tensor& l) : data(l.data) {}
		sparse_tensor(sparse_tensor&& l) noexcept : data(std::move(l.data)) {}
		sparse_tensor& operator=(const sparse_tensor& l) { data = l.data; return *this; }
		sparse_tensor& operator=(sparse_tensor&& l) noexcept { data = std::move(l.data); return *this; }

		// for the i-th column, return the indices
		index_p index(size_t i) { return data.colptr + i * rank(); }
		const_index_p index(size_t i) const { return data.colptr + i * rank(); }
		T& val(size_t i) { return data.valptr[i]; }
		const T& val(size_t i) const { return data.valptr[i]; }

		index_v index_vector(size_t i) const {
			index_v result(rank());
			for (size_t j = 0; j < rank(); j++)
				result[j] = index(i)[j];
			return result;
		}

		inline size_t alloc() const { return data.alloc; }
		inline size_t nnz() const { return data.rowptr[1]; }
		inline size_t rank() const { return data.rank - 1; }
		inline std::vector<size_t> dims() const {
			std::vector<size_t> result(data.dims.begin() + 1, data.dims.end());
			return result;
		}
		inline size_t dim(size_t i) const { return data.dims[i + 1]; }
		inline bool check_sorted() const { return data.check_sorted(); }
		inline void zero() { data.zero(); }
		inline void reserve(size_t size) { data.reserve(size); }
		inline void shrink_to_fit() { data.shrink_to_fit(); }
		inline void resize(size_t new_nnz) {
			if (new_nnz > alloc())
				reserve(new_nnz);
			data.rowptr[1] = new_nnz;
		}

		// we assume that the tensor is sorted
		inline std::vector<size_t> rowptr() const {
			std::vector<size_t> result(dim(0) + 1);
			result[0] = 0;
			for (size_t i = 0; i < nnz(); i++) {
				result[index(i)[0] + 1]++;
			}
			for (size_t i = 0; i < dim(0); i++)
				result[i + 1] += result[i];
			return result;
		}

		// change the dimensions of the tensor
		// it is dangerous, only for internal use
		inline void change_dims(const std::vector<size_t>& new_dims) {
			data.change_dims(prepend_num(new_dims, (size_t)1));
		}

		inline void flatten(const std::vector<std::vector<size_t>>& pos) {
			auto r = rank();
			auto nr = pos.size();
			std::vector<index_t> newindex(nr);
			std::vector<size_t> new_dims(nr);
			auto old_dim = dims();
			auto init_ptr = data.colptr;

			// first compute new dimensions
			for (size_t i = 0; i < nr; i++) {
				new_dims[i] = 1;
				for (auto j : pos[i])
					new_dims[i] *= old_dim[j];
			}
			new_dims = prepend_num(new_dims, (size_t)1);

			for (size_t i = 0; i < nnz(); i++) {
				auto ptr = index(i);
				for (size_t j = 0; j < nr; j++) {
					newindex[j] = 0;
					for (auto k : pos[j])
						newindex[j] = newindex[j] * old_dim[k] + ptr[k];
				}
				for (size_t j = 0; j < nr; j++)
					init_ptr[i * nr + j] = newindex[j];
			}
			// colptr must stay able to hold alloc * (rank - 1) = nr * alloc entries, because the
			// copy constructor reads exactly that many; sizing the buffer by nnz and then shrinking
			// alloc to nnz (with reserve) throws the spare capacity away and makes the next
			// push_back reallocate and grow again
			if (data.alloc != 0)
				data.colptr = s_realloc(data.colptr, nr * data.alloc);

			// change the dimensions
			data.dims = new_dims;
			data.rank = nr + 1;
		}

		// The mixed radix conversion of reshape(): the flat index of an entry under old_dims, split
		// again by new_dims. Acc is uint64_t whenever the product of the dimensions fits in 64 bits
		// -- every shape that is not astronomically large -- and int_t (a big integer) otherwise,
		// so the arithmetic is only as wide as it has to be. One implementation, two widths.
		template <typename Acc>
		void reshape_impl(const std::vector<size_t>& old_dims, const std::vector<size_t>& new_dims, index_t* newcolptr) const {
			const auto r = rank();
			const auto m = new_dims.size();
			for (size_t i = 0; i < nnz(); i++) {
				auto ptr = index(i);
				Acc flat = 0;
				for (size_t j = 0; j < r; j++)
					flat = flat * old_dims[j] + ptr[j];
				for (size_t j = m; j > 0; j--) {
					const Acc tmp = flat % new_dims[j - 1];
					flat /= new_dims[j - 1];
					if constexpr (std::is_same_v<Acc, int_t>)
						newcolptr[i * m + j - 1] = (index_t)(tmp.to_si());
					else
						newcolptr[i * m + j - 1] = (index_t)(tmp);
				}
			}
		}

		// reshape, for example {2,100} to {2,5,20}
		inline void reshape(const std::vector<size_t>& new_dims) {
			auto old_dims = dims();
			if (new_dims == old_dims)
				return; // the identity: the coordinates already are the ones of new_dims

			// same invariant as in flatten: alloc * new_dims.size() entries, not nnz * ...
			index_t* newcolptr = (data.alloc == 0) ? nullptr : s_malloc<index_t>(data.alloc * new_dims.size());

			// the flat index is bounded by the product of the dimensions, so a shape whose product
			// leaves 64 bits is the only one that needs the big integer path
			uint64_t prod = 1;
			bool fits_u64 = true;
			for (const auto d : old_dims) {
				if (d != 0 && prod > std::numeric_limits<uint64_t>::max() / d) {
					fits_u64 = false;
					break;
				}
				prod *= d;
			}

			if (fits_u64)
				reshape_impl<uint64_t>(old_dims, new_dims, newcolptr);
			else
				reshape_impl<int_t>(old_dims, new_dims, newcolptr);

			s_free(data.colptr);
			data.colptr = newcolptr;
			data.dims = prepend_num(new_dims, (size_t)1);
			data.rank = new_dims.size() + 1;
		}

		inline void insert(const index_v& l, const T& val, bool mode = true) { data.insert(prepend_num(l), val, mode); }
		inline void insert_add(const index_v& l, const T& val) { data.insert_add(prepend_num(l), val); }
		void push_back(const_index_p l, const T& new_val) {
			auto n_nnz = nnz();
			if (n_nnz + 1 > data.alloc)
				reserve((data.alloc + 1) * 2);
			s_copy(index(n_nnz), l, rank());
			val(n_nnz) = new_val;
			data.rowptr[1]++; // increase the nnz
		}
		void push_back(const index_v& l, const T& new_val) { push_back(l.data(), new_val); }
		void pop_back() { if (data.rowptr[1] != 0) data.rowptr[1]--; }
		inline void canonicalize() { data.canonicalize(); }
		inline void sort_indices(thread_pool* pool = nullptr) { data.sort_indices(pool); }
		inline sparse_tensor transpose(const std::vector<size_t>& perm, thread_pool* pool = nullptr, const bool sort_ind = true) const {
			std::vector<size_t> perm_new(perm);
			for (auto& a : perm_new) { a++; }
			perm_new = prepend_num(perm_new, (size_t)0);
			sparse_tensor B;
			B.data = data.transpose(perm_new, pool, sort_ind);
			B.sort_indices(pool);
			return B;
		}

		// take a span of elements
		// sparse_tensor.take(levelspec, {start, end}) returns a sparse_tensor whose elements have the levelspec-th index in range [start, end)
		// elements in the resulting sparse_tensor have their levelspec-th indices reindexed in [0, end - start)
		sparse_tensor take(const size_t levelspec, const std::pair<size_t, size_t>& span, thread_pool* pool = nullptr) const {
			sparse_tensor B;
			B.data = data.take(levelspec + 1, span, pool);
			return B;
		}

		// extract a (rank-1) tensor by fixing the levelspec-th index
		sparse_tensor extract(const size_t levelspec, const size_t index, thread_pool* pool = nullptr) const {
			sparse_tensor B;
			B.data = data.extract(levelspec + 1, index, pool);
			return B;
		}

		/* extract certain part from tensor, similar to Part in Mathematica 
		   part specification for each level:
		   - `index_t`: extract that index
		   - `std::vector<std::pair<index_t, index_t>>`: take the union of spans `[first, second)`
		   Note: the result must be itself a sparse_tensor, i.e., has rank >= 2.
		*/
		sparse_tensor part(const std::vector<std::variant<index_t, std::pair<index_t, index_t>>>& partspec, thread_pool* pool = nullptr) const {
			std::vector<std::variant<index_t, std::pair<index_t, index_t>>> newpartspec;
			newpartspec.push_back(std::pair{0, 1}); // for the first dimension
			newpartspec.insert(newpartspec.end(), partspec.begin(), partspec.end());
			sparse_tensor B;
			B.data = data.part(newpartspec, pool);
			return B;
		}

		sparse_mat<T, index_t> to_sparse_mat(thread_pool* pool = nullptr, const bool sort_ind = true) const & {
			if (rank() != 2) {
				report_error("sparse_tensor.to_sparse_mat: rank must be 2");
				return sparse_mat<T, index_t>();
			}
			// sorting cannot be done in place on a const tensor: sort a copy and convert that
			if (sort_ind && !check_sorted()) {
				sparse_tensor sorted(*this);
				sorted.sort_indices(pool);
				return sorted.to_sparse_mat(pool, false);
			}
			auto r = dim(0);
			auto c = dim(1);
			auto rptr = rowptr();

			sparse_mat<T, index_t> mat(r, c);
			// one row body, run by the pool or in a plain loop
			auto copy_row = [&](const size_t i) {
				const size_t nz = rptr[i + 1] - rptr[i];
				mat[i].reserve(nz);
				mat[i].resize(nz);
				if (nz == 0)
					return;
				std::copy(data.valptr + rptr[i], data.valptr + rptr[i + 1], mat[i].entries);
				// skip the first index, which is the row index
				auto ptr = index(rptr[i]) + 1;
				for (size_t j = 0; j < nz; j++)
					mat[i].indices[j] = ptr[2 * j];
				};
			if (pool == nullptr) {
				for (size_t i = 0; i < r; i++)
					copy_row(i);
			}
			else {
				pool->detach_loop(0, r, copy_row);
				pool->wait();
			}
			return mat;
		}

		// true when the entries already are in the order that gen_perm_by() would put them in, i.e.
		// when the permutation it returns is the identity: a caller that only walks the entries in that
		// order can ask this and skip building the permutation (one size_t per entry) altogether
		bool is_ordered_by(const std::vector<size_t>* order = nullptr) const {
			// a positions list of the wrong size is ignored by gen_perm, which orders by the natural
			// order instead, so it is ignored here as well
			if (order != nullptr && order->size() != rank())
				order = nullptr;
			const auto r = rank();
			const auto nz = nnz();
			if (order == nullptr) {
				for (size_t i = 1; i < nz; i++)
					if (!(lexico_compare(index(i - 1), index(i), r) <= 0))
						return false;
			}
			else {
				for (size_t i = 1; i < nz; i++)
					if (!(lexico_compare(index(i - 1), index(i), *order) <= 0))
						return false;
			}
			return true;
		}

		// the same conversion for a tensor that is about to be discarded: the values are moved out of
		// it, and it is left empty
		sparse_mat<T, index_t> to_sparse_mat(thread_pool* pool = nullptr, const bool sort_ind = true) && {
			if (rank() != 2) {
				report_error("sparse_tensor.to_sparse_mat: rank must be 2");
				return sparse_mat<T, index_t>();
			}

			sparse_mat<T, index_t> mat(dim(0), dim(1));
			std::vector<size_t> rownnz(dim(0) + 1, 0);
			bool rows_sorted = true;
			for (size_t i = 0; i < nnz(); i++) {
				rownnz[index(i)[0] + 1]++;
				if (i != 0 && index(i - 1)[0] > index(i)[0])
					rows_sorted = false;
			}

			if (pool != nullptr && rows_sorted) {
				// turn row nnz into offset
				for (size_t i = 0; i < mat.nrow; i++)
					rownnz[i + 1] += rownnz[i];
				pool->detach_loop(0, mat.nrow, [&](size_t i) {
					auto nz = rownnz[i + 1] - rownnz[i];
					if (nz == 0)
						return;
					mat[i].reserve(nz);
					mat[i].resize(nz);
					for (size_t j = 0; j < nz; j++)
						mat[i].indices[j] = index(rownnz[i] + j)[1];
					std::move(data.valptr + rownnz[i], data.valptr + rownnz[i + 1], mat[i].entries);
					});
				pool->wait();
			}
			else {
				for (size_t i = 0; i < mat.nrow; i++)
					mat[i].reserve(rownnz[i + 1]);
				for (size_t i = 0; i < nnz(); i++) {
					auto ptr = index(i);
					mat[ptr[0]].push_back(ptr[1], std::move(val(i)));
				}
			}
			clear();

			if (sort_ind) {
				if (pool == nullptr) {
					for (size_t i = 0; i < mat.nrow; i++)
						mat[i].sort_indices();
				}
				else {
					pool->detach_loop(0, mat.nrow, [&](size_t i) { mat[i].sort_indices(); });
					pool->wait();
				}
			}
			return mat;
		}

		// the permutation that orders the entries by the labels at the positions listed in order, or by
		// all the labels in the natural order when order is nullptr
		// the entries are usually already in that order, because the callers hand in tensors that were
		// built sorted or returned by an operation that sorts: verifying it costs one cheap pass, so
		// the identity permutation is returned as soon as the check succeeds
		std::vector<size_t> gen_perm_by(const std::vector<size_t>* order, thread_pool* pool) const {
			const auto r = rank();
			const auto nz = nnz();
			// the counting passes below keep one bucket per label, so only the orderings whose
			// dimensions add up to a manageable number of buckets are counted

			std::vector<size_t> perm = perm_init(nz);
			if (is_ordered_by(order))
				return perm;

			// the ordering can also be obtained by counting the labels one position at a time: a
			// label is smaller than the dimension of its position, so each counting pass only
			// needs a bucket per label, which is much cheaper than comparing tuples when the
			// dimensions are small; the positions are counted from the last to the first, so that
			// the first one decides the order, and one pass is stable, so the passes compose
			const size_t len = order == nullptr ? r : order->size();
			size_t buckets = 0;
			bool countable = len > 0 && len <= r;
			for (size_t l = 0; l < len && countable; l++) {
				const size_t d = dim(order == nullptr ? l : (*order)[l]);
				countable = d > 0 && d <= counting_max_buckets - buckets;
				buckets += d;
			}

			// the ordering by comparisons, which is what the counting passes fall back on when they
			// cannot be used or when they meet an entry that steps outside its dimensions
			auto by_comparison = [&]() {
				auto less = [&](size_t a, size_t b) {
					return order == nullptr ? lexico_compare(index(a), index(b), r) < 0
						: lexico_compare(index(a), index(b), *order) < 0;
					};
				parallel_sort(perm.begin(), perm.end(), less, pool);
				return perm;
				};

			if (countable) {
				std::vector<index_t> label(nz);
				std::vector<size_t> next(nz);
				std::vector<size_t> count;
				for (size_t l = len; l-- > 0;) {
					const size_t pos = order == nullptr ? l : (*order)[l];
					const size_t d = dim(pos);

					count.assign(d, 0);
					// the labels are read through the permutation, which walks the index array at
					// random; the pass stays serial: it is one read per entry, and the counting pass
					// below it is serial anyway, so a dispatch would cost more than the pass
					for (size_t i = 0; i < nz; i++)
						label[i] = index(perm[i])[pos];

					for (size_t i = 0; i < nz; i++) {
						const auto value = label[i];
						if (value < 0 || static_cast<size_t>(value) >= d)
							return by_comparison();
						count[value]++;
					}

					size_t start = 0;
					for (size_t b = 0; b < d; b++) {
						const size_t c = count[b];
						count[b] = start;
						start += c;
					}
					for (size_t i = 0; i < nz; i++)
						next[count[label[i]]++] = perm[i];
					perm.swap(next);
				}
				return perm;
			}

			return by_comparison();
		}

		std::vector<size_t> gen_perm(thread_pool* pool = nullptr) const { return gen_perm_by(nullptr, pool); }

		std::vector<size_t> gen_perm(const std::vector<size_t>& index_perm, thread_pool* pool = nullptr) const {
			// index_perm is a permutation of the positions of the index vector, so it must have
			// exactly rank() entries; the caller is expected to check this, we only degrade to the
			// natural order here instead of handing back a perm of the wrong size
			if (index_perm.size() != rank()) {
				report_error("gen_perm: index_perm size is not equal to rank");
				return gen_perm(pool);
			}

			return gen_perm_by(&index_perm, pool);
		}

		void transpose_replace(const std::vector<size_t>& perm, thread_pool* pool = nullptr, const bool sort_ind = true) {
			// perm[i] is the old position of the new i-th index, so perm has to be a permutation of
			// [0, rank()); anything else would read out of bounds below
			std::vector<bool> used(rank(), false);
			bool valid = perm.size() == rank();
			for (size_t i = 0; valid && i < perm.size(); i++) {
				if (perm[i] >= rank() || used[perm[i]])
					valid = false;
				else
					used[perm[i]] = true;
			}
			if (!valid) {
				report_error("transpose_replace: perm must be a permutation of the indices");
				return;
			}

			std::vector<size_t> new_dims(rank() + 1);
			new_dims[0] = data.dims[0];

			for (size_t i = 0; i < rank(); i++)
				new_dims[i + 1] = data.dims[perm[i] + 1];
			data.dims = new_dims;

			auto method = [&](size_t ss, size_t ee) {
				std::vector<size_t> index_new(rank());
				for (size_t i = ss; i < ee; i++) {
					auto ptr = index(i);
					for (size_t j = 0; j < rank(); j++)
						index_new[j] = ptr[perm[j]];
					std::copy(index_new.begin(), index_new.end(), ptr);
				}
				};

			if (pool == nullptr)
				method(0, nnz());
			else {
				pool->detach_blocks(0, nnz(), method);
				pool->wait();
			}
			if (sort_ind)
				sort_indices(pool);
		}

		sparse_tensor<T, index_t, SPARSE_COO> chop(size_t pos, size_t aa) const {
			std::vector<size_t> dims_new = dims();
			dims_new.erase(dims_new.begin() + pos);
			sparse_tensor<T, index_t, SPARSE_COO> result(dims_new);
			index_v index_new;
			index_new.reserve(rank() - 1);
			for (size_t i = 0; i < nnz(); i++) {
				if (index(i)[pos] != aa)
					continue;
				for (size_t j = 0; j < rank(); j++) {
					if (j != pos)
						index_new.push_back(index(i)[j]);
				}
				result.push_back(index_new, val(i));
				index_new.clear();
			}
			return result;
		}

		sparse_tensor& operator=(const sparse_tensor<T, index_t, SPARSE_CSR>& l) {
			if (alloc() == 0) {
				init(l.dims(), l.nnz());
			}
			else {
				change_dims(l.dims());
				reserve(l.nnz());
			}
			resize(l.nnz());

			auto r = rank();
			auto n_row = dim(0);

			// first copy the data
			s_copy(data.valptr, l.data.valptr, l.nnz());

			// then deal with the indices
			for (size_t i = 0; i < n_row; i++) {
				for (size_t j = l.data.rowptr[i]; j < l.data.rowptr[i + 1]; j++) {
					auto tmp_index = index(j);
					tmp_index[0] = i;
					for (size_t k = 0; k < r - 1; k++)
						tmp_index[k + 1] = l.data.colptr[j * (r - 1) + k];
				}
			}

			return *this;
		}

		sparse_tensor(const sparse_tensor<T, index_t, SPARSE_CSR>& l) { *this = l; }

		// The stolen data is in the CSR layout; turn it into the COO one: the row index goes in
		// front of every entry, rowptr becomes {0, nnz} (which frees the CSR row pointer array),
		// and dims gets the artificial leading 1. The stride is widened in place, backwards, so
		// that no second tensor-sized index buffer is needed (the source of an entry can only be
		// overlapped by that entry's own row slot, which is written last).
		void csr_to_coo_layout() noexcept {
			auto r = data.rank;
			auto n_row = data.dims[0];

			if (data.alloc == 0) {
				s_free(data.colptr); // no capacity: nothing to widen, only the buffer to drop
				data.colptr = nullptr;
			}
			else {
				data.colptr = s_realloc<index_t>(data.colptr, data.alloc * r);
			}
			auto colptr = data.colptr;
			for (size_t i = n_row; i-- > 0;) {
				for (size_t j = data.rowptr[i + 1]; j-- > data.rowptr[i];) {
					auto oldptr = colptr + j * (r - 1);
					auto newptr = colptr + j * r;
					for (size_t k = r - 1; k-- > 0;)
						newptr[k + 1] = oldptr[k];
					newptr[0] = i;
				}
			}

			data.rowptr = { 0, data.rowptr.back() };
			data.dims = prepend_num(data.dims, (size_t)1);
			data.rank++;
		}

		sparse_tensor& operator=(sparse_tensor<T, index_t, SPARSE_CSR>&& l) noexcept {
			data = std::move(l.data);
			csr_to_coo_layout();

			return *this;
		}

		sparse_tensor(sparse_tensor<T, index_t, SPARSE_CSR>&& l) noexcept {
			data = std::move(l.data);
			csr_to_coo_layout();
		}

		sparse_tensor(const sparse_mat<T, index_t>& mat, thread_pool* pool = nullptr) {
			*this = sparse_tensor<T, index_t, SPARSE_CSR>(mat, pool);
		}

		sparse_tensor& operator=(const sparse_mat<T, index_t>& mat) {
			*this = sparse_tensor<T, index_t, SPARSE_CSR>(mat);
			return *this;
		}

		sparse_tensor(sparse_mat<T, index_t>&& mat, thread_pool* pool = nullptr) {
			*this = sparse_tensor<T, index_t, SPARSE_CSR>(std::move(mat), pool);
		}

		sparse_tensor& operator=(sparse_mat<T, index_t>&& mat) {
			*this = sparse_tensor<T, index_t, SPARSE_CSR>(std::move(mat));
			return *this;
		}

		void print_test() const {
			for (size_t j = 0; j < data.rowptr[1]; j++) {
				for (size_t k = 0; k < data.rank - 1; k++)
					std::cout << (size_t)(data.colptr[j * (data.rank - 1) + k]) << " ";
				std::cout << " : " << data.valptr[j] << std::endl;
			}
		}
	};

	// some other functions

	// join two sparse matrices
	template <typename T, typename index_t>
	sparse_mat<T, index_t> sparse_mat_join(const sparse_mat<T, index_t>& A, const sparse_mat<T, index_t>& B, thread_pool* pool = nullptr) {
		sparse_mat<T, index_t> res(A.nrow + B.nrow, std::max(A.ncol, B.ncol));

		// one row body, run by the pool or in a plain loop: the same in both branches
		auto append_rows = [&](const sparse_mat<T, index_t>& src, const size_t offset) {
			auto copy_row = [&](const size_t i) { res[i + offset] = src[i]; };
			if (pool == nullptr) {
				for (size_t i = 0; i < src.nrow; i++)
					copy_row(i);
			}
			else {
				pool->detach_loop(0, src.nrow, copy_row);
				pool->wait();
			}
			};
		append_rows(A, 0);
		append_rows(B, A.nrow);

		return res;
	}

	template <typename T, typename index_t>
	sparse_mat<T, index_t> sparse_mat_join(sparse_mat<T, index_t>&& A, sparse_mat<T, index_t>&& B) {
		sparse_mat<T, index_t> res = std::move(A);
		res.append(std::move(B));
		res.ncol = std::max(A.ncol, B.ncol);
		return res;
	}

	// split a sparse matrix into two parts
	template <typename T, typename index_t>
	std::pair<sparse_mat<T, index_t>, sparse_mat<T, index_t>> sparse_mat_split(const sparse_mat<T, index_t>& mat, const size_t split_row, thread_pool* pool = nullptr) {
		if (split_row > mat.nrow) {
			report_error("sparse_mat_split: split_row out of range");
			return {};
		}

		sparse_mat<T, index_t> A(split_row, mat.ncol);
		sparse_mat<T, index_t> B(mat.nrow - split_row, mat.ncol);

		auto split_into = [&](sparse_mat<T, index_t>& dst, const size_t first) {
			const size_t n = dst.nrow;
			auto copy_row = [&](const size_t i) { dst[i] = mat[first + i]; };
			if (pool == nullptr) {
				for (size_t i = 0; i < n; i++)
					copy_row(i);
			}
			else {
				pool->detach_loop(0, n, copy_row);
				pool->wait();
			}
			};
		split_into(A, 0);
		split_into(B, split_row);

		return { A, B };
	}

	template <typename T, typename index_t>
	std::pair<sparse_mat<T, index_t>, sparse_mat<T, index_t>> sparse_mat_split(sparse_mat<T, index_t>&& mat, const size_t split_row) {
		if (split_row > mat.nrow) {
			report_error("sparse_mat_split: split_row out of range");
			return {};
		}

		sparse_mat<T, index_t> A(split_row, mat.ncol);
		sparse_mat<T, index_t> B(mat.nrow - split_row, mat.ncol);

		A.rows = std::vector<sparse_vec<T, index_t>>(std::make_move_iterator(mat.rows.begin()),
			std::make_move_iterator(mat.rows.begin() + split_row));
		B.rows = std::vector<sparse_vec<T, index_t>>(std::make_move_iterator(mat.rows.begin() + split_row),
			std::make_move_iterator(mat.rows.end()));

		mat.clear();

		return { A, B };
	}

	// a submatrix view of sparse_mat
	// it only contains a pointer to the original matrix and a list of row indices
	// and it does not own any memory
	// to save memory, if rows[0] > mat.nrow, then it is a full view
	template <typename T, typename index_t>
	struct sparse_mat_subview {
		const sparse_mat<T, index_t>* mat_ptr = nullptr;
		std::vector<size_t> rows;

		sparse_mat_subview() = default;
		~sparse_mat_subview() = default;

		sparse_mat_subview(sparse_mat<T, index_t>& mat_) {
			mat_ptr = &mat_;
			rows = { mat_.nrow + 1 }; // full view
		}

		sparse_mat_subview(const sparse_mat<T, index_t>& mat_) {
			mat_ptr = &mat_;
			rows = { mat_.nrow + 1 }; // full view
		}

		sparse_mat_subview(sparse_mat<T, index_t>& mat_, const std::vector<size_t>& rows_) {
			mat_ptr = &mat_;
			rows = rows_;
		}

		sparse_mat_subview(const sparse_mat_subview& l) { mat_ptr = l.mat_ptr; rows = l.rows; }
		sparse_mat_subview(sparse_mat_subview&& l) noexcept { mat_ptr = l.mat_ptr; rows = std::move(l.rows); }
		sparse_mat_subview& operator=(const sparse_mat_subview& l) {
			if (this == &l)
				return *this;
			mat_ptr = l.mat_ptr;
			rows = l.rows;
			return *this;
		}
		sparse_mat_subview& operator=(sparse_mat_subview&& l) noexcept {
			if (this == &l)
				return *this;
			mat_ptr = l.mat_ptr;
			rows = std::move(l.rows);
			return *this;
		}

		size_t nrow() const {
			if (mat_ptr) {
				if (rows[0] > mat_ptr->nrow) // full view
					return mat_ptr->nrow;
				return rows.size();
			}
			return 0;
		}

		size_t ncol() const {
			if (mat_ptr)
				return mat_ptr->ncol;
			return 0;
		}

		size_t nnz() const {
			if (mat_ptr) {
				size_t result = 0;
				if (rows[0] > mat_ptr->nrow) // full view
					return mat_ptr->nnz();
				for (auto r : rows)
					result += (*mat_ptr)[r].nnz();
				return result;
			}
			return 0;
		}

		bool is_full() const {
			if (mat_ptr) {
				return rows[0] > mat_ptr->nrow; // full view
			}
			return false;
		}

		// the index of the i-th row in the underlying matrix, it is dangerous because we do
		// not check the index
		size_t row_of(const size_t i) const {
			if (rows[0] > mat_ptr->nrow) // full view
				return i;
			return rows[i];
		}

		// access the i-th row of the subview, it is dangerous because we do not check the index
		sparse_vec<T, index_t>& operator[](size_t i) {
			return const_cast<sparse_mat<T, index_t>&>(*mat_ptr)[row_of(i)];
		}
		const sparse_vec<T, index_t>& operator[](size_t i) const {
			return (*mat_ptr)[row_of(i)];
		}

		const size_t operator()(size_t i) const {
			return row_of(i);
		}

		sparse_mat<T, index_t>& get_mat() {
			return const_cast<sparse_mat<T, index_t>&>(*mat_ptr);
		}
		const sparse_mat<T, index_t>& get_mat() const {
			return *mat_ptr;
		}

		void traverse(std::function<void(size_t)> func) {
			if (mat_ptr == nullptr)
				return;
			if (rows[0] > mat_ptr->nrow) { // full view
				for (size_t i = 0; i < mat_ptr->nrow; i++)
					func(i);
			}
			else {
				for (auto r : rows)
					func(r);
			}
		}
	};

}

#endif