Restructured a small bit of code.