Change branch checking logic

@VladimirMikhailov, it seems well to me! I'll add this code.
Thanks