HMSegmentControl扩展红点功能

HMSegmentControl没有红点功能，在不修改pod源码情况下，如何通过扩展给其增加红点？

image-20201020175938450.png

首先分析分段控制器实现（HMSegmentedControl.m）

- (void)drawRect:(CGRect)rect {
    // Remove all sublayers to avoid drawing images over existing ones
    self.scrollView.layer.sublayers = nil;
        ...
        
    CATextLayer *titleLayer = [CATextLayer layer];
    titleLayer.frame = rect;
    titleLayer.alignmentMode = kCAAlignmentCenter;
    if ([UIDevice currentDevice].systemVersion.floatValue < 10.0 ) {
        titleLayer.truncationMode = kCATruncationEnd;
    }
    titleLayer.string = [self attributedTitleAtIndex:idx];
    titleLayer.contentsScale = [[UIScreen mainScreen] scale];
    [self.scrollView.layer addSublayer:titleLayer];
}

发现其用drawRect实现，要么写一个子类，在里边增加代码；要么用运行时来处理。我选择后者

源码可供参考：分类（HMSegmentedControl+Ext）

HMSegmentedControl+Ext.h
//留给外部赋值用
- (void)setBadgeIndex:(NSInteger) index;

HMSegmentedControl+Ext.m
+ (void)load{
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken,^{
        Class cls = [HMSegmentedControl class];
        LYSwizzleInstanceMethod(cls, @selector(drawRect:), cls, @selector(my_drawRect:));
    });
}

- (void)my_drawRect:(CGRect)rect{
    [self my_drawRect:rect];
    [self updateBadge];
}

//设置红点的下标，并用关联属性记录下来
- (void)setBadgeIndex:(NSInteger) index{
    objc_setAssociatedObject(self, @"seg_badge", @(index), OBJC_ASSOCIATION_RETAIN);
}

- (void)updateBadge{
    NSNumber *obj = objc_getAssociatedObject(self, @"seg_badge");
    NSInteger index = obj.integerValue;
    if (!self.scrollView) {
        return;
    }
    
    NSArray <CALayer *> *layers = [self.scrollView.layer sublayers];
    if (layers.count == 0 || layers.count <= index) {
        return;
    }
    
    NSMutableArray <CALayer *> *marray = [NSMutableArray array];
    for (CALayer *layer in layers) {
        if ([layer isKindOfClass:CATextLayer.class]) {
            [marray addObject:layer];
        }
    }
    if (marray.count == 0 || marray.count <= index) {
        return;
    }
    
    //查找红点所对应的layer,并初始化一个红点在上面
    CALayer *superLayer = marray[index];
    CALayer *redLayer = [CALayer layer];
    redLayer.frame = CGRectMake(superLayer.frame.size.width-8-3, 0, 8, 8);
    redLayer.backgroundColor = [UIColor redColor].CGColor;
    redLayer.cornerRadius = 4.0f;
    redLayer.masksToBounds = YES;
    [superLayer addSublayer:redLayer];
}