Expandable group for groupie that takes a size of elements to show for the collapsed state

ExpandableGroup.kt

import androidx.recyclerview.widget.DiffUtil
import androidx.recyclerview.widget.ListUpdateCallback
import com.xwray.groupie.Group
import com.xwray.groupie.Item
import com.xwray.groupie.NestedGroup
import com.xwray.groupie.ViewHolder

// Inspired from ExpandableGroup
// https://github.com/lisawray/groupie/blob/master/library/src/main/java/com/xwray/groupie/ExpandableGroup.java
class ExpandableGroup(
    sizeCollapsed: Int = 0,
    isExpanded: Boolean = false
) : NestedGroup() {

    var sizeCollapsed = sizeCollapsed
        private set
    var isExpanded = isExpanded
        private set
    private val children = ArrayList<Group>()

    fun setSizeCollapsed(size: Int) {
        if (size >= 0) {
            sizeCollapsed = size
            if (size == 0) {
                notifyChanged()
            } else {
                notifyItemRangeChanged(0, size)
            }
        }
    }

    fun setExpanded(expanded: Boolean) {
        val oldSize = itemCount
        isExpanded = expanded
        val newSize = itemCount
        if (oldSize > newSize) {
            notifyItemRangeRemoved(newSize, oldSize - newSize)
        } else {
            notifyItemRangeInserted(oldSize, newSize - oldSize)
        }
    }

    fun toggleExpanded() {
        setExpanded(!isExpanded)
    }

    fun clear() {
        if (itemCount > 0) {
            update(emptyList())
        }
    }

    private fun getPositions(position: Int, count: Int): Pair<Int, Int>? {
        return when {
            isExpanded -> position to count
            position < sizeCollapsed -> {
                position to Math.min(sizeCollapsed - position - 1, count)
            }
            else -> null
        }
    }

    // Taken from com.xwray.groupie.Section
    private val listUpdateCallback = object : ListUpdateCallback {

        override fun onInserted(position: Int, count: Int) {
            val pair = getPositions(position, count)
            if (pair != null) {
                notifyItemRangeInserted(pair.first, pair.second)
            }
        }

        override fun onRemoved(position: Int, count: Int) {
            val pair = getPositions(position, count)
            if (pair != null) {
                notifyItemRangeRemoved(pair.first, pair.second)
            }
        }

        override fun onMoved(fromPosition: Int, toPosition: Int) {
            val pair = getPositions(fromPosition, toPosition)
            if (pair != null) {
                notifyItemMoved(pair.first, pair.second)
            }
        }

        override fun onChanged(position: Int, count: Int, payload: Any?) {
            val pair = getPositions(position, count)
            if (pair != null) {
                notifyItemRangeChanged(pair.first, pair.second, payload)
            }
        }
    }

    fun update(newBodyGroups: Collection<Group>) {
        val oldBodyGroups = ArrayList(children)
        val oldBodyItemCount = getItemCount(oldBodyGroups)
        val newBodyItemCount = getItemCount(newBodyGroups)

        val diffResult = DiffUtil.calculateDiff(object : DiffUtil.Callback() {
            override fun getOldListSize(): Int {
                return oldBodyItemCount
            }

            override fun getNewListSize(): Int {
                return newBodyItemCount
            }

            override fun areItemsTheSame(oldItemPosition: Int, newItemPosition: Int): Boolean {
                val oldItem = getItem(oldBodyGroups, oldItemPosition)
                val newItem = getItem(newBodyGroups, newItemPosition)
                return newItem.isSameAs(oldItem)
            }

            override fun areContentsTheSame(oldItemPosition: Int, newItemPosition: Int): Boolean {
                val oldItem = getItem(oldBodyGroups, oldItemPosition)
                val newItem = getItem(newBodyGroups, newItemPosition)
                return newItem == oldItem
            }

            override fun getChangePayload(oldItemPosition: Int, newItemPosition: Int): Any? {
                val oldItem = getItem(oldBodyGroups, oldItemPosition)
                val newItem = getItem(newBodyGroups, newItemPosition)
                return oldItem.getChangePayload(newItem)
            }
        })

        super.removeAll(children)
        children.clear()
        children.addAll(newBodyGroups)
        super.addAll(newBodyGroups)

        diffResult.dispatchUpdatesTo(listUpdateCallback)
    }

    override fun add(group: Group) {
        super.add(group)
        if (isExpanded) {
            val itemCount = itemCount
            children.add(group)
            notifyItemRangeInserted(itemCount, group.itemCount)
        } else {
            children.add(group)
        }
    }

    override fun addAll(groups: Collection<Group>) {
        if (groups.isEmpty()) {
            return
        }
        super.addAll(groups)

        val itemCount = itemCount

        if (isExpanded) {
            this.children.addAll(groups)
            notifyItemRangeInserted(itemCount, getItemCount(groups))
        } else {
            this.children.addAll(groups)
            val itemsInGroups = getItemCount(groups)
            if (itemCount < sizeCollapsed) {
                val itemsToAdd = Math.min(sizeCollapsed - itemCount, itemsInGroups)
                notifyItemRangeInserted(itemCount, itemsToAdd)
            }
        }
    }

    override fun addAll(position: Int, groups: Collection<Group>) {
        if (groups.isEmpty()) {
            return
        }
        super.addAll(position, groups)
        val itemCount = itemCount
        if (isExpanded) {
            this.children.addAll(position, groups)
            notifyItemRangeInserted(itemCount, getItemCount(groups))
        } else {
            this.children.addAll(position, groups)
            if (itemCount < sizeCollapsed) {
                val itemsToAdd = sizeCollapsed - itemCount
                notifyItemRangeInserted(itemCount, itemsToAdd)
            }
        }
    }

    override fun remove(group: Group) {
        if (!this.children.contains(group)) return
        super.remove(group)
        val position = getItemCountBeforeGroup(group)
        if (isExpanded) {
            children.remove(group)
            notifyItemRangeRemoved(position, group.itemCount)
        } else {
            children.remove(group)
            if (position < sizeCollapsed) {
                val itemsToAdd = sizeCollapsed - position
                notifyItemRangeRemoved(position, itemsToAdd)
            }
        }
    }

    override fun removeAll(groups: Collection<Group>) {
        if (groups.isEmpty() || !this.children.containsAll(groups)) return
        super.removeAll(groups)
        if (isExpanded) {
            this.children.removeAll(groups)
            for (group in groups) {
                val position = getItemCountBeforeGroup(group)
                children.remove(group)
                notifyItemRangeRemoved(position, group.itemCount)
            }
        } else {
            for (group in groups) {
                val position = getItemCountBeforeGroup(group)
                children.remove(group)
                if (position < sizeCollapsed) {
                    val itemsToRemove = sizeCollapsed - position
                    notifyItemRangeRemoved(position, itemsToRemove)
                }
            }
        }
    }

    override fun getGroup(position: Int): Group {
        return children[position]
    }

    override fun getPosition(group: Group): Int {
        return children.indexOf(group)
    }

    override fun getGroupCount(): Int {
        return if (isExpanded) {
            children.size
        } else {
            Math.min(children.size, sizeCollapsed)
        }
    }

    private fun dispatchChildChanges(group: Group): Boolean {
        val groupPosition = getPosition(group)
        return isExpanded || groupPosition in 0..(sizeCollapsed - 1)
    }

    override fun onChanged(group: Group) {
        if (dispatchChildChanges(group)) {
            super.onChanged(group)
        }
    }

    override fun onItemInserted(group: Group, position: Int) {
        if (dispatchChildChanges(group)) {
            super.onItemInserted(group, position)
        }
    }

    override fun onItemChanged(group: Group, position: Int) {
        if (dispatchChildChanges(group)) {
            super.onItemChanged(group, position)
        }
    }

    override fun onItemChanged(group: Group, position: Int, payload: Any) {
        if (dispatchChildChanges(group)) {
            super.onItemChanged(group, position, payload)
        }
    }

    override fun onItemRemoved(group: Group, position: Int) {
        if (dispatchChildChanges(group)) {
            super.onItemRemoved(group, position)
        }
    }

    override fun onItemRangeChanged(group: Group, positionStart: Int, itemCount: Int) {
        if (dispatchChildChanges(group)) {
            super.onItemRangeChanged(group, positionStart, itemCount)
        }
    }

    override fun onItemRangeChanged(group: Group, positionStart: Int, itemCount: Int, payload: Any) {
        if (dispatchChildChanges(group)) {
            super.onItemRangeChanged(group, positionStart, itemCount, payload)
        }
    }

    override fun onItemRangeInserted(group: Group, positionStart: Int, itemCount: Int) {
        if (dispatchChildChanges(group)) {
            super.onItemRangeInserted(group, positionStart, itemCount)
        }
    }

    override fun onItemRangeRemoved(group: Group, positionStart: Int, itemCount: Int) {
        if (dispatchChildChanges(group)) {
            super.onItemRangeRemoved(group, positionStart, itemCount)
        }
    }

    override fun onItemMoved(group: Group, fromPosition: Int, toPosition: Int) {
        if (dispatchChildChanges(group)) {
            super.onItemMoved(group, fromPosition, toPosition)
        }
    }

    companion object {

        private fun getItem(groups: Collection<Group>, position: Int): Item<ViewHolder> {
            var previousPosition = 0

            groups.forEach { group ->
                val size = group.itemCount
                if (size + previousPosition > position) {
                    return group.getItem(position - previousPosition)
                }
                previousPosition += size
            }

            throw IndexOutOfBoundsException(
                "Wanted item at " + position + " but there are only "
                        + previousPosition + " items"
            )
        }
    }
}