Small code rework