[IMP] Improve code bases on PR review