Modified code structure a bit...