ios - How to merge two different sections in collection view

I want to merge two different sections that uses two different cells respectively(cellA and cellB). Using the following code I could get a collection view with two sections.

import UIKit

class ViewController: UIViewController, UICollectionViewDelegate, UICollectionViewDataSource{

    @IBOutlet var testCollectionView: UICollectionView!
    
    override func viewDidLoad() {
        super.viewDidLoad()
        self.testCollectionView.delegate = self
        self.testCollectionView.dataSource = self
    }
   
    func numberOfSections(in collectionView: UICollectionView) -> Int {
        return 2
    }
    
    func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
        if section == 0 {
            return 3
        } else {
            return 1
        }
    }
    
    func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
        if indexPath.section == 0 {
            let cell = collectionView.dequeueReusableCell(withReuseIdentifier: "taskCell", for: indexPath) as! CellA
            cell.layer.cornerRadius = 2
            cell.layer.shadowColor = UIColor(red:0.82, green:0.82, blue:0.82, alpha:1.0).cgColor
            cell.layer.shadowOffset = CGSize.zero
            cell.layer.shadowOpacity = 1
            cell.layer.shadowRadius = 4
            cell.layer.masksToBounds = false
            return cell
        } else {
            let cell2 = collectionView.dequeueReusableCell(withReuseIdentifier: "addNewCell", for: indexPath) as! CellB
            return cell2
        }
    }
}

extension ViewController: UICollectionViewDelegateFlowLayout {
    func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, sizeForItemAt indexPath: IndexPath) -> CGSize {
        
        let width = testCollectionView.bounds.width/2 - 30
        let height = width
        
        return CGSize(width: width, height: height)
    }
    func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, insetForSectionAt section: Int) -> UIEdgeInsets {
        return UIEdgeInsetsMake(30, 20, 10, 20)
    }
    func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, minimumLineSpacingForSectionAt section: Int) -> CGFloat {
        return 20
    }
    func collectionView(_ collectionView: UICollectionView, didSelectItemAt indexPath: IndexPath) {
        print(indexPath.row)
    }
}

And the output looks like this.

Now, I want the section 1(CellB) to be merged with section 0(CellA). Like this.

I couldn't find the solution to achieve this. Any help will be appreciated.

You can show both cell types in the same section with these changes:

1. Add properties for the number of aCells and bCells. This is better than putting the magic numbers 3 and 1 in your code because it documents what the numbers represent. If you ever need to change them, you can change them in one spot.
2. Change numberOfSections to return 1.
3. Change numberOfItemsInSection to return the sum of the number of A cells and the number of B cells: return aCells + bCells.
4. In cellForItemAt, compare indexPath.item to aCells to figure out when to switch over to B cells.

var aCells = 3
var bCells = 1

func numberOfSections(in collectionView: UICollectionView) -> Int {
    return 1
}

func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
    return aCells + bCells
}

func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
    if indexPath.item < aCells {
        let cell = collectionView.dequeueReusableCell(withReuseIdentifier: "taskCell", for: indexPath) as! CellA
        cell.layer.cornerRadius = 2
        cell.layer.shadowColor = UIColor(red:0.82, green:0.82, blue:0.82, alpha:1.0).cgColor
        cell.layer.shadowOffset = CGSize.zero
        cell.layer.shadowOpacity = 1
        cell.layer.shadowRadius = 4
        cell.layer.masksToBounds = false
        return cell
    } else {
        let cell2 = collectionView.dequeueReusableCell(withReuseIdentifier: "addNewCell", for: indexPath) as! CellB
        return cell2
    }
}