simplify the code logic and reduce memory a little bit Signed-off-by: Xiang Xiao <xiaoxiang@xiaomi.com>