/*************************************************************************
 * Copyright (c) 2018-2020, NVIDIA CORPORATION. All rights reserved.
 *
 * See LICENSE.txt for license information
 ************************************************************************/

#ifndef SCCL_CPUSET_H_
#define SCCL_CPUSET_H_

#include "base.h"

namespace sccl {
namespace hardware {
namespace topology {
namespace topo {
// Convert local_cpus, e.g. 0003ff,f0003fff to cpu_set_t

/**
 * 将十六进制字符转换为对应的整数值
 *
 * @param c 输入的十六进制字符（0-9, a-f）
 * @return 返回对应的整数值（0-15），如果输入无效则返回-1
 */
static int hexToInt(char c) {
    int v = c - '0';
    if(v < 0)
        return -1;
    if(v > 9)
        v = 10 + c - 'a';
    if((v < 0) || (v > 15))
        return -1;
    return v;
}

#define CPU_SET_N_U32 (sizeof(cpu_set_t) / sizeof(uint32_t))

/**
 * 将十六进制字符串转换为CPU集合掩码
 *
 * @param str 输入的十六进制字符串，用逗号分隔不同部分
 * @param mask 输出的CPU集合掩码
 * @return scclSuccess 表示转换成功
 *
 * @note 字符串从左到右对应掩码从高到低的32位字
 *       每个字符代表4位十六进制数
 *       遇到非十六进制字符会提前终止转换
 */
static scclResult_t scclStrToCpuset(const char* str, cpu_set_t* mask) {
    uint32_t cpumasks[CPU_SET_N_U32];
    int m       = CPU_SET_N_U32 - 1;
    cpumasks[m] = 0;
    for(int o = 0; o < strlen(str); o++) {
        char c = str[o];
        if(c == ',') {
            m--;
            cpumasks[m] = 0;
        } else {
            int v = hexToInt(c);
            if(v == -1)
                break;
            cpumasks[m] <<= 4;
            cpumasks[m] += v;
        }
    }
    // Copy cpumasks to mask
    for(int a = 0; m < CPU_SET_N_U32; a++, m++) {
        memcpy(((uint32_t*)mask) + a, cpumasks + m, sizeof(uint32_t));
    }
    return scclSuccess;
}

/**
 * 将CPU集合掩码转换为十六进制字符串表示
 *
 * @param mask 输入的CPU集合掩码
 * @param str 输出的字符串缓冲区，用于存储转换结果
 * @return 返回操作结果(scclSuccess表示成功)
 *
 * 转换规则：
 * 1. 将cpu_set_t按字节从高到低转换为十六进制字符串
 * 2. 每4个字节后添加一个逗号分隔符
 * 3. 忽略前导零
 */
static scclResult_t scclCpusetToStr(cpu_set_t* mask, char* str) {
    int c       = 0;
    uint8_t* m8 = (uint8_t*)mask;
    for(int o = sizeof(cpu_set_t) - 1; o >= 0; o--) {
        if(c == 0 && m8[o] == 0)
            continue;
        sprintf(str + c, "%02x", m8[o]);
        c += 2;
        if(o && o % 4 == 0) {
            sprintf(str + c, ",");
            c++;
        }
    }
    str[c] = '\0';
    return scclSuccess;
}

/**
 * 将CPU集合掩码转换为范围字符串表示
 *
 * @param mask 输入的CPU集合掩码
 * @param str  用于存储结果的缓冲区
 * @param len  缓冲区长度
 * @return     返回转换后的字符串指针(即str参数)
 *
 * 该函数将CPU集合掩码转换为可读的范围字符串格式，例如"0-3,5,7-9"。
 * 如果缓冲区空间不足，结果会被截断。空集合会返回空字符串。
 */
static char* scclCpusetToRangeStr(cpu_set_t* mask, char* str, size_t len) {
    int c     = 0;
    int start = -1;
    // Iterate through all possible CPU bits plus one extra position
    for(int cpu = 0; cpu <= CPU_SETSIZE; cpu++) {
        int isSet = (cpu == CPU_SETSIZE) ? 0 : CPU_ISSET(cpu, mask);
        // Start of a new range
        if(isSet && start == -1) {
            start = cpu;
        }
        // End of a range, add comma between ranges
        if(!isSet && start != -1) {
            if(cpu - 1 == start) {
                c += snprintf(str + c, len - c, "%s%d", c ? "," : "", start);
            } else {
                c += snprintf(str + c, len - c, "%s%d-%d", c ? "," : "", start, cpu - 1);
            }
            if(c >= len - 1)
                break;
            start = -1;
        }
    }
    if(c == 0)
        str[0] = '\0';
    return str;
}

} // namespace topo
} // namespace topology
} // namespace hardware
} // namespace sccl
#endif