I'm attaching a small patch with the changes I suggested above.