Minor code structure adjustment